Kaynağa Gözat

Reformatting

Kirk Trombley 3 yıl önce
ebeveyn
işleme
d5a4312a49
3 değiştirilmiş dosya ile 131 ekleme ve 101 silme
  1. 37 32
      convert.py
  2. 73 63
      ingest.py
  3. 21 6
      nearest.py

+ 37 - 32
convert.py

@@ -1,22 +1,27 @@
 # I could not find a single decent RGB -> CIELUV conversion library out there
 
+
 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
+    # 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],
+    [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
@@ -31,25 +36,25 @@ 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)
+    # 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)

+ 73 - 63
ingest.py

@@ -7,80 +7,90 @@ from convert import rgb_to_cieluv
 
 
 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)
+    # returns true if a pixel is transparent or pure black
+    return a == 0 or (r, g, b) == (0, 0, 0)
 
 
 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)
+    # 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)
 
 
 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))
+    # 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"])
+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,
-  )
+    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([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")
+    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")

+ 21 - 6
nearest.py

@@ -12,13 +12,24 @@ def norm(r, g, b):
 
 
 parser = ArgumentParser()
-parser.add_argument("color", nargs="?", default=None, help="Target color, randomized if not provided")
-parser.add_argument("-n", "--number", default=1, type=int, help="Number of Pokemon to find")
-parser.add_argument("-c", "--closeness", default=2, type=int, help="Closeness coefficient")
+parser.add_argument(
+    "color", nargs="?", default=None, help="Target color, randomized if not provided"
+)
+parser.add_argument(
+    "-n", "--number", default=1, type=int, help="Number of Pokemon to find"
+)
+parser.add_argument(
+    "-c", "--closeness", default=2, type=int, help="Closeness coefficient"
+)
 parser.add_argument("-d", "--database", default="database.csv", help="Database file")
 parser.add_argument("-x", "--exclude-x", action="store_true", help="Exclude X")
 parser.add_argument("-z", "--normalize", action="store_true", help="Normalize q and Y")
-parser.add_argument("-l", "--convert-luv", action="store_true", help="Convert input color to CIELUV before calculation")
+parser.add_argument(
+    "-l",
+    "--convert-luv",
+    action="store_true",
+    help="Convert input color to CIELUV before calculation",
+)
 parser.add_argument("-v", "--verbose", action="store_true", help="Print raw scores")
 args = parser.parse_args()
 
@@ -29,7 +40,11 @@ if args.color is not None:
     cleaned_color = args.color.strip("#")
     if len(cleaned_color) != 6:
         raise ValueError("Color must be a 6 digit hex")
-    color = (int(cleaned_color[0:2], base=16), int(cleaned_color[2:4], base=16), int(cleaned_color[4:6], base=16))
+    color = (
+        int(cleaned_color[0:2], base=16),
+        int(cleaned_color[2:4], base=16),
+        int(cleaned_color[4:6], base=16),
+    )
 else:
     color = tuple(int(random.random() * 255) for _ in range(3))
     print(f"Generated random color: #{''.join(hex(c)[2:] for c in color)} / {color}")
@@ -53,7 +68,7 @@ with open(args.database) as infile:
         results.append((score, name))
 
 if args.number > 1:
-    for i, (score, name) in enumerate(sorted(results)[:args.number]):
+    for i, (score, name) in enumerate(sorted(results)[: args.number]):
         print(f"{i+1}:\t{name}")
         if args.verbose:
             print(f"\t\t\t{score=}")