فهرست منبع

Rewrite ingester to also calculate CIELUV values

Kirk Trombley 3 سال پیش
والد
کامیت
67b2265775
1فایلهای تغییر یافته به همراه127 افزوده شده و 33 حذف شده
  1. 127 33
      ingest.py

+ 127 - 33
ingest.py

@@ -1,45 +1,139 @@
 #!/usr/bin/env python3
+from collections import namedtuple
+
 from PIL import Image
 
 
-def ingest_png(file_name: str) -> tuple[str, int, int, int, int]:
-    print(f"Ingesting {file_name}")
+def rescale_and_linearize(component: int) -> float:
+  # takes an sRGB color component [0,255]
+  # first rescales to [0,1]
+  # then linearizes according to some CIEXYZ stuff I don't understand
+  # then rescales to [0, 100]
+  component /= 255
+  linearized = component / 12.92 if component <= 0.04045 else ((component + 0.055) / 1.055) ** 2.4
+  return 100 * linearized
+
+
+# conversion values I also do not understand
+# pulled from https://www.image-engineering.de/library/technotes/958-how-to-convert-between-srgb-and-ciexyz
+# instead of easy rgb, since it seemed to give more accurate values
+rgb_to_xyz_matrix = [
+  [0.4124564, 0.3575761, 0.1804375],
+  [0.2126729, 0.7151522, 0.0721750],
+  [0.0193339, 0.1191920, 0.9503041],
+]
+
+# reference values I also also do not understand
+# pulled from easy rgb
+# note X and Y here have nothing to do with the X and Y metrics below
+ref_x = 95.047
+ref_y = 100.000
+ref_z = 108.883
+ref_denom = ref_x + 15 * ref_y + 3 * ref_z
+ref_u = 4 * ref_x / ref_denom
+ref_v = 9 * ref_y / ref_denom
+
+
+def rgb_to_cieluv(r: int, g: int, b: int) -> tuple[float, float, float]:
+  # accepts RGB (components [0, 255])
+  # converts to CIE LUV (components [0, 1])
+  # math taken from http://www.easyrgb.com/en/math.php
+  
+  # RGB (components [0, 255]) -> XYZ (components [0, 100])
+  # X, Y and Z output refer to a D65/2° standard illuminant.
+  sr, sg, sb = (rescale_and_linearize(c) for c in (r, g, b))
+  x, y, z = (cr * sr + cg * sg + cb * sb for cr, cg, cb in rgb_to_xyz_matrix)
+
+  # XYZ (components [0, 100]) -> LUV (components [0, 100])
+  uv_denom = x + 15 * y + 3 * z
+  u = 4 * x / uv_denom
+  v = 9 * y / uv_denom
+
+  if y > 0.8856:
+    yprime = (y / 100) ** (1/3)
+  else:
+    yprime = (y / 100) * 7.787 + (16/116)
+
+  lstar = 116 * yprime - 16
+  lstar_factor = 13 * lstar
+  return lstar, lstar_factor * (u - ref_u), lstar_factor * (v - ref_v)
+
 
-    # image name - strip leading path and trailing extension
-    name = file_name.rsplit("/", maxsplit=1)[1].split(".", maxsplit=1)[0]
+def is_outline(r: int, g: int, b: int, a: int) -> bool:
+  # returns true if a pixel is transparent or pure black
+  return a == 0 or (r, g, b) == (0, 0, 0)
 
-    # read non-transparent pixels of image as RGB values
-    pixels = [(r, g, b) for r, g, b, a in Image.open(file_name).convert("RGBA").getdata() if a > 0 and (r > 0 or g > 0 or b > 0)]
 
-    # X metric - mean sq norm of all pixels
-    x = sum(x * x + y * y + z * z for x, y, z in pixels) / len(pixels)
+def x_metric(pixels: list[tuple[float, float, float]]) -> float:
+  # X metric - the mean squared Euclidean norm
+  # computed as the sum of the squares of the components of the pixels,
+  # normalized by the number of pixels
+  return sum(comp * comp for pix in pixels for comp in pix) / len(pixels)
 
-    # Y metrics - mean of color components
-    yr_sum = 0
-    yg_sum = 0
-    yb_sum = 0
-    for r, g, b in pixels:
-        yr_sum += r
-        yg_sum += g
-        yb_sum += b
-    yr = yr_sum / len(pixels)
-    yg = yg_sum / len(pixels)
-    yb = yb_sum / len(pixels)
 
-    return name, x, yr, yg, yb
+def y_metric(pixels: list[tuple[float, float, float]]) -> tuple[float, float, float]:
+  # Y metric - the mean pixel of the image
+  return tuple(sum(p[i] for p in pixels) / len(pixels) for i in range(3))
+
+
+ImageInfo = namedtuple("ImageInfo", ["name", "xrgb", "xluv", "yr", "yg", "yb", "yl", "yu", "yv"])
+
+
+def ingest_png(file_name: str) -> ImageInfo:
+  print(f"Ingesting {file_name}")
+
+  # image name - strip leading path and trailing extension
+  name = file_name.rsplit("/", maxsplit=1)[1].split(".", maxsplit=1)[0]
+
+  # read non-outline pixels of image 
+  rgb_pixels = [(r, g, b) 
+            for r, g, b, a in Image.open(file_name).convert("RGBA").getdata() 
+            if not is_outline(r, g, b, a)]
+  
+  # convert RGB pixels to CIELUV values
+  luv_pixels = [rgb_to_cieluv(*p) for p in rgb_pixels]
+
+  # compute and return metrics
+  xrgb = x_metric(rgb_pixels)
+  xluv = x_metric(luv_pixels)
+  yr, yg, yb = y_metric(rgb_pixels)
+  yl, yu, yv = y_metric(luv_pixels)
+  return ImageInfo(
+    name=name,
+    xrgb=xrgb,
+    xluv=xluv,
+    yr=yr,
+    yg=yg,
+    yb=yb,
+    yl=yl,
+    yu=yu,
+    yv=yv,
+  )
 
 
 if __name__ == "__main__":
-    import csv
-    import os
-
-    data = [ingest_png("pngs/" + fn) for f in os.listdir("pngs") if (fn := os.fsdecode(f)).endswith(".png")]
-    with open("database.csv", "w") as outfile:
-        writer = csv.writer(outfile, delimiter=",", quotechar="'")
-        writer.writerows(data)
-
-    with open("database.js", "w") as outfile:
-        outfile.write("const database = [\n")
-        for name, x, yr, yg, yb in data:
-            outfile.write(f'  [ "{name}", {x}, {yr}, {yg}, {yb} ],\n')
-        outfile.write("];\n")
+  import csv
+  import os
+
+  data = [ingest_png("pngs/" + fn) for f in os.listdir("pngs") if (fn := os.fsdecode(f)).endswith(".png")]
+  
+  with open("database.csv", "w") as outfile:
+    writer = csv.writer(outfile, delimiter=",", quotechar="'")
+    writer.writerows([d.name, d.xrgb, d.yr, d.yg, d.yb] for d in data)
+
+  with open("database-luv.csv", "w") as outfile:
+    writer = csv.writer(outfile, delimiter=",", quotechar="'")
+    writer.writerows([d.name, d.xluv, d.yl, d.yu, d.yv] for d in data)
+
+  with open("database.js", "w") as outfile:
+    outfile.write("const database = [\n")
+    for info in data:
+      fields = ", ".join((
+        f'name: "{info.name}"',
+        f"xRGB: {info.xrgb}",
+        f"xLUV: {info.xluv}",
+        f"yRGB: [ {info.yr}, {info.yg}, {info.yb} ]",
+        f"yLUV: [ {info.yl}, {info.yu}, {info.yv} ]",
+      ))
+      outfile.write(f"  {{ {fields} }},\n")
+    outfile.write("];\n")