ingest.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. from re import M
  2. from PIL import Image
  3. def calc_X_metric(pixels: list[tuple[int, int, int]]) -> float:
  4. return sum(x * x + y * y + z * z for x, y, z in pixels) / len(pixels)
  5. def calc_Y_metric(pixels: list[tuple[int, int, int]]) -> tuple[float, float, float]:
  6. r_sum = 0
  7. g_sum = 0
  8. b_sum = 0
  9. for r, g, b in pixels:
  10. r_sum += r
  11. g_sum += g
  12. b_sum += b
  13. factor = 2 / len(pixels)
  14. return (r_sum * factor, g_sum * factor, b_sum * factor)
  15. def read_png_data(img_name: str) -> list[tuple[int, int, int]]:
  16. return [(r, g, b) for r, g, b, a in Image.open(img_name).convert("RGBA").getdata() if a > 0]
  17. def ingest_png(file_name: str):
  18. print(f"Ingesting {file_name}")
  19. data = read_png_data(file_name)
  20. x = calc_X_metric(data)
  21. y = calc_Y_metric(data)
  22. name = file_name.rsplit("/", maxsplit=1)[1].split(".", maxsplit=1)[0]
  23. return name, x, *y
  24. if __name__ == "__main__":
  25. import csv
  26. import os
  27. data = [ingest_png("pngs/" + fn) for f in os.listdir("pngs") if (fn := os.fsdecode(f)).endswith(".png")]
  28. with open("database.csv", "w") as outfile:
  29. writer = csv.writer(outfile, delimiter=",", quotechar="'")
  30. writer.writerows(data)