Add override engine

Author: janek-cosmose

add_panel_override.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python3
+import hashlib
+import json
+from pathlib import Path
+
+import typer
+from PIL import Image
+
+app = typer.Typer()
+
+
+def add_panel_override(image_path: Path, panel_no: list[int]) -> None:
+    overrides = json.load(open('parse_qwantz/data/panel_overrides.json'))
+    md5 = hashlib.md5(Image.open(image_path).tobytes()).hexdigest()
+    if md5 not in overrides:
+        comic_id, file_name = image_path.name.split(' - ')
+        overrides[md5] = {"comic_id": int(comic_id), "file_name": file_name, "panels": {}}
+    for n in panel_no:
+        overrides[md5]["panels"][n] = ["", ""]
+    json.dump(overrides, open('parse_qwantz/data/panel_overrides.json', 'w'), indent=2, ensure_ascii=False)
+
+
+if __name__ == '__main__':
+    typer.run(add_panel_override)

parse_qwantz/main.py

@@ -1,3 +1,4 @@
+import hashlib
 import logging
 import sys
 from dataclasses import dataclass
@@ -11,6 +12,7 @@
 from parse_qwantz.color_logs import set_logging_formatter, ColorFormatter
 from parse_qwantz.colors import Color
 from parse_qwantz.lines import Line
+from parse_qwantz.panel_overrides import get_panel_overrides
 from parse_qwantz.text_blocks import get_text_blocks, TextBlock
 from parse_qwantz.elements import get_elements
 from parse_qwantz.match_blocks import match_blocks
@@ -74,8 +76,16 @@ def __bool__(self):
 
 
 def parse_qwantz(image: Image, debug: bool, log_to_file: bool) -> Iterable[list[str]]:
-    masked = prepare_image(image)
+    md5 = hashlib.md5(image.tobytes()).hexdigest()
+    panel_overrides = get_panel_overrides().get(md5, {})
+    masked, good_panels = prepare_image(image)
     for i, (panel, characters) in enumerate(zip(PANELS, CHARACTERS), start=1):
+        if str(i) in panel_overrides:
+            yield panel_overrides[str(i)]
+            continue
+        if i not in good_panels:
+            logger.warning("Non-standard panel without an override")
+            continue
         set_current_panel(i, log_to_file)
         (width, height), (x, y) = panel
         cropped = masked.crop((x, y, x + width, y + height))

parse_qwantz/panel_overrides.py

@@ -0,0 +1,17 @@
+import json
+from dataclasses import dataclass
+from functools import cache
+from importlib.resources import files, as_file
+
+import parse_qwantz
+
+OVERRIDE_FILE_PATH = files(parse_qwantz).joinpath('data/panel_overrides.json')
+
+
+@cache
+def get_panel_overrides() -> dict[str, dict[int, list[str]]]:
+    with as_file(OVERRIDE_FILE_PATH) as override_path:
+        return {
+            key: value["panels"]
+            for key, value in json.load(open(override_path)).items()
+        }

parse_qwantz/prepare_image.py

@@ -16,10 +16,11 @@
 
 SAMPLE = [
     ((113, 183), (128, 255, 64)),
+    ((290, 160), (255, 128, 161)),
     ((704, 183), (255, 242, 179)),
+    ((143, 375), (255, 191, 82)),
     ((452, 405), (255, 191, 82)),
-    ((290, 160), (255, 128, 161)),
-    ((372, 484), (0, 0, 0)),
+    ((585, 386), (69, 196, 5)),
 ]
 
 
@@ -28,22 +29,27 @@ class ImageError(Exception):
 
 
 @cache
-def get_mask_image():
+def get_mask_image() -> Image:
     with as_file(MASK_FILE_PATH) as image_path:
         return Image.open(image_path)
 
 
-def prepare_image(image: Image):
+def prepare_image(image: Image) -> tuple[Image, list[int]]:
     if image.size != DIM:
         logger.error(f"Wrong image dimensions: {image.size}, only {DIM} is valid")
         raise ImageError(f"Wrong image dimensions: {image.size}, only {DIM} is valid")
     palette = image.getpalette()
     palette = tuple(palette) if palette else None
-    for pixel, expected_color in SAMPLE:
+    good_panels: list[int] = []
+    for panel_no, (pixel, expected_color) in enumerate(SAMPLE, start=1):
         color = image.getpixel(pixel)
         color = normalize_color(color, palette)
-        if square_distance(color, expected_color) > COLOR_THRESHOLD:
-            logger.error(f"Invalid template: expected {expected_color} at {pixel}; found {color}")
-            raise ImageError(f"Invalid template: expected {expected_color} at {pixel}; found {color}")
+        if square_distance(color, expected_color) <= COLOR_THRESHOLD:
+            good_panels.append(panel_no)
+        else:
+            logger.warning(f"Invalid template: expected {expected_color} at {pixel}; found {color}")
+    if not good_panels:
+        logger.error("Invalid template")
+        raise ImageError(f"Invalid template")
     all_white = Image.new(mode='RGB', size=DIM, color=(255, 255, 255))
-    return Image.composite(image, all_white, get_mask_image())
+    return Image.composite(image, all_white, get_mask_image()), good_panels

sort_panel_overrides.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+import hashlib
+import json
+from pathlib import Path
+
+import typer
+from PIL import Image
+
+app = typer.Typer()
+
+
+def sort_panel_overrides() -> None:
+    overrides = json.load(open('parse_qwantz/data/panel_overrides.json'))
+    sorted_overrides = {key: value for key, value in sorted(overrides.items(), key=lambda r: r[1]["comic_id"])}
+    json.dump(sorted_overrides, open('parse_qwantz/data/panel_overrides.json', 'w'), indent=2, ensure_ascii=False)
+
+
+if __name__ == '__main__':
+    typer.run(sort_panel_overrides)