Merge pull request #48 from timotheeg/fix_palette_compute_script

Capture areas from nes logical pixels, and update easiercap palette accordingly

nestris_ocr/palettes/easiercap.json

@@ -1 +1 @@
-[[[64, 38, 233], [83, 168, 255]], [[0, 128, 0], [124, 200, 1]], [[141, 9, 184], [213, 93, 255]], [[63, 39, 230], [81, 212, 38]], [[157, 16, 88], [59, 207, 108]], [[61, 208, 109], [133, 130, 255]], [[149, 36, 0], [89, 90, 86]], [[99, 14, 220], [84, 0, 17]], [[64, 39, 232], [153, 41, 0]], [[150, 37, 0], [214, 136, 26]]]
+[[[65, 40, 228], [86, 168, 254]], [[3, 127, 1], [125, 199, 6]], [[139, 14, 179], [211, 96, 255]], [[64, 41, 225], [85, 210, 45]], [[154, 19, 87], [64, 205, 111]], [[65, 206, 112], [134, 131, 255]], [[146, 38, 2], [90, 91, 87]], [[97, 18, 215], [82, 0, 20]], [[64, 40, 227], [151, 43, 4]], [[147, 39, 3], [211, 137, 32]]]

scripts/compute_color_palette.py

@@ -1,23 +1,45 @@
 # invoke from root as python3 -m scripts.compute_color_palette easiercap nestris_ocr/assets/sample_inputs/easiercap/lvl%d/field.png
 import json
 import sys
-from PIL import Image, ImageFilter
+from PIL import Image
 import numpy as np
 from nestris_ocr.colors import Colors
 from math import sqrt
 
 
+def getAverageColor(pixels):
+    return (
+        round(sqrt(sum([p[0] * p[0] for p in pixels]) / len(pixels))),
+        round(sqrt(sum([p[1] * p[1] for p in pixels]) / len(pixels))),
+        round(sqrt(sum([p[2] * p[2] for p in pixels]) / len(pixels))),
+    )
+
+
 def getColors(level, field):
     img = Image.open(field)
 
-    # extra smoothness from PIL
-    slight_blur = ImageFilter.GaussianBlur()
-    slight_blur.radius = 1
+    res_img = Image.new("RGBA", (img.width, img.height), (0, 0, 0, 255),)
 
-    img = img.filter(slight_blur)
+    res_img.paste(img.convert("RGBA"), (img.width * 0, 0))
 
-    spanx = img.width / 10
-    spany = img.height / 20
+    # nes pix size, block span and capture are only valid on perfect calibration
+    # perfect calibration implies capturing the right and bottom black edges
+    # still, since we're capturing a safe zone, things should still work with small discrepancies
+
+    nes_pixels_to_sample = (
+        (2, 4.5),
+        (4.5, 2),
+        (4.5, 4.5),
+    )
+
+    nes_pix_xsize = img.width / 80
+    nes_pix_ysize = img.height / 160
+
+    spanx = nes_pix_xsize * 8
+    spany = nes_pix_ysize * 8
+
+    corner_highlight = Image.new("RGBA", (1, 1), (255, 0, 0, 128))
+    capture_highlight = Image.new("RGBA", (1, 1), (0, 255, 0, 128))
 
     np_img = np.array(img, dtype=np.uint16)
 
@@ -28,41 +50,43 @@ def getColors(level, field):
 
     for y in range(20):
         for x in range(10):
-            xidx = round(spanx * (x + 0.5))
-            yidx = round(spany * (y + 0.5))
+            xidx = spanx * x
+            yidx = spany * y
+
+            # print the top-left corner
+            res_img.paste(
+                corner_highlight, (round(xidx), round(yidx)), corner_highlight
+            )
 
-            pix = [0, 0, 0]
+            pixels = []
 
-            # grab 9 pixels in a 3x3 square for each block
-            # on the lower-right quadrant
-            # and compute color average
-            for i in range(xidx, xidx + 3):
-                for j in range(yidx, yidx + 3):
-                    tmp = np_img[j, i]
-                    pix[0] += tmp[0] * tmp[0]
-                    pix[1] += tmp[1] * tmp[1]
-                    pix[2] += tmp[2] * tmp[2]
+            for offsetx, offsety in nes_pixels_to_sample:
+                right = round(xidx + nes_pix_xsize * offsetx)
+                top = round(yidx + nes_pix_ysize * offsety)
+                left = round(xidx + nes_pix_xsize * (offsetx + 1))
+                bottom = round(yidx + nes_pix_ysize * (offsety + 1))
 
-            pix[0] = sqrt(pix[0] / 9)
-            pix[1] = sqrt(pix[1] / 9)
-            pix[2] = sqrt(pix[2] / 9)
+                cap = capture_highlight.resize((left - right, bottom - top))
+
+                res_img.paste(cap, (right, top), cap)
+
+                # grab all pixels in the capture area
+                for i in range(right, left):
+                    for j in range(top, bottom):
+                        pixels.append(np_img[j, i])
+
+            pix = getAverageColor(pixels)
 
             col_index = colors.getClosestColorIndex(pix)
 
-            color_instances[col_index].append(pix)
+            color_instances[col_index].extend(pixels)
 
-    # average the colors
-    col1 = (
-        round(sum([c[0] for c in color_instances[2]]) / len(color_instances[2])),
-        round(sum([c[1] for c in color_instances[2]]) / len(color_instances[2])),
-        round(sum([c[2] for c in color_instances[2]]) / len(color_instances[2])),
-    )
+    if level == 0:
+        res_img.convert("RGB").resize((img.width * 3, img.height * 3)).show()
 
-    col2 = (
-        round(sum([c[0] for c in color_instances[3]]) / len(color_instances[3])),
-        round(sum([c[1] for c in color_instances[3]]) / len(color_instances[3])),
-        round(sum([c[2] for c in color_instances[3]]) / len(color_instances[3])),
-    )
+    # average the colors
+    col1 = getAverageColor(color_instances[2])
+    col2 = getAverageColor(color_instances[3])
 
     return col1, col2
 
@@ -127,6 +151,7 @@ def getPaletteImage(palette):
 
     palette_img.save("nestris_ocr/palettes/%s.png" % (palette_name,))
 
-    # show both reference palette and computed palette for comparison
     palette_img.show()
-    getPaletteImage(colors.palette).show()
+
+    # shows the default palette for comparison
+    # getPaletteImage(colors.palette).show()