Update app.py

app.py

@@ -1,82 +1,148 @@
+#!/usr/bin/env python
+"""
+Kovaaks PNG Metadata Cloner — v2
+↓ Improvements
+   • enforces PNG chunk order (IHDR → sRGB → gAMA → pHYs → eXIf → (…) → IDAT → IEND)
+   • keeps only a single instance of each ancillary chunk (no duplicates)
+   • validates / auto-converts target to 50 × 50 RGBA if the user ticks a checkbox
+   • warns when resulting file > 256 B (Kovaaks soft limit)
+"""
+
+from __future__ import annotations
+
 import os
 import struct
 import tempfile
+from pathlib import Path
 from typing import BinaryIO, Iterator
 
 import gradio as gr
+from PIL import Image
 
 PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
-
-
-# ────────────────────────────────────────────────────────── helpers ──
+# ancillary chunks Kovaaks cares about, in desired order
+ANCILLARY_ORDER = (
+    b"sRGB",
+    b"gAMA",
+    b"pHYs",
+    b"eXIf",  # empty container in template crosshair
+    b"iTXt",
+    b"tEXt",
+    b"zTXt",
+)
+
+# ───────────────────────────────────────────────── helpers ──
 def _read_chunks(fp: BinaryIO) -> Iterator[tuple[bytes, bytes]]:
     """
     Yield (raw_chunk_bytes, chunk_type_bytes) for every chunk in *fp*.
-    `raw_chunk_bytes` includes length, type, data and CRC exactly as in file,
-    so we can copy them verbatim later.
-    The PNG signature must already be consumed by the caller.
+    Caller must have consumed the 8-byte PNG signature already.
     """
     while True:
         length_bytes = fp.read(4)
         if not length_bytes:
             break  # EOF
         if len(length_bytes) != 4:
-            raise ValueError("Corrupted PNG: unexpected EOF while reading chunk length.")
+            raise ValueError("Corrupted PNG: could not read length")
 
-        length = struct.unpack(">I", length_bytes)[0]          # ← fixed (no stray space)
+        length = struct.unpack(">I", length_bytes)[0]
         chunk_type = fp.read(4)
-        data       = fp.read(length)
-        crc        = fp.read(4)
+        data = fp.read(length)
+        crc = fp.read(4)
 
         if len(chunk_type) != 4 or len(data) != length or len(crc) != 4:
-            raise ValueError("Corrupted PNG: truncated chunk.")
+            raise ValueError("Corrupted PNG: truncated chunk")
 
         yield length_bytes + chunk_type + data + crc, chunk_type
 
 
-def clone_metadata(src_path: str, tgt_path: str, out_path: str) -> str:
-    """
-    Copy all ancillary chunks (anything except IHDR / IDAT / IEND) from *src_path*
-    into *tgt_path* immediately after its IHDR, then save to *out_path*.
-    Pixel data of the target stays untouched.
-    """
+def _clean_ancillary(src_chunks: list[tuple[bytes, bytes]]) -> list[bytes]:
+    """Return at most one instance of each ordered ancillary chunk."""
+    kept: dict[bytes, bytes] = {}
+    for raw, tp in src_chunks:
+        if tp in ANCILLARY_ORDER and tp not in kept:
+            kept[tp] = raw
+    # return in canonical order
+    return [kept[tag] for tag in ANCILLARY_ORDER if tag in kept]
+
+
+def _validate_or_resize(path: str, auto_fix: bool) -> None:
+    """Ensure 50×50 RGBA; optionally fix in-place (overwrites file)."""
+    img = Image.open(path)
+    if img.size == (50, 50) and img.mode == "RGBA":
+        return  # already good
+
+    if not auto_fix:
+        raise ValueError("Target must be 50×50 px, RGBA (enable auto-fix to convert).")
+
+    # auto-convert
+    img = img.convert("RGBA")
+    img = img.resize((50, 50), Image.LANCZOS)
+    img.save(path, optimize=True)
+
+
+def clone_metadata(
+    src_path: str, tgt_path: str, out_path: str, auto_fix: bool = True
+) -> str:
+    """Clone cleaned ancillary chunks from src → tgt. Save to out_path."""
     with open(src_path, "rb") as fs, open(tgt_path, "rb") as ft:
         if fs.read(8) != PNG_SIGNATURE or ft.read(8) != PNG_SIGNATURE:
-            raise ValueError("Both files must be valid PNG images.")
+            raise ValueError("Both uploads must be valid PNG images.")
 
         src_chunks = list(_read_chunks(fs))
         tgt_chunks = list(_read_chunks(ft))
 
-    # take every non-critical chunk from source
-    ancillary = [
-        raw for raw, tp in src_chunks
-        if tp not in (b"IHDR", b"IDAT", b"IEND")
-    ]
+    ancillary = _clean_ancillary(src_chunks)
+
+    # Remove from target any ancillary tags we will replace
+    replace_tags = {tp for _, tp in ancillary}
+    rebuilt: list[bytes] = []
 
-    out_chunks: list[bytes] = []
+    inserted = False
     for raw, tp in tgt_chunks:
-        out_chunks.append(raw)
-        if tp == b"IHDR":                     # inject right after IHDR
-            out_chunks.extend(ancillary)
+        if tp == b"IHDR":
+            rebuilt.append(raw)
+            # insert ordered ancillary exactly once
+            if not inserted:
+                rebuilt.extend(ancillary)
+                inserted = True
+            continue
+
+        if tp in replace_tags:
+            # skip existing duplicate in target
+            continue
+
+        rebuilt.append(raw)
 
     with open(out_path, "wb") as fo:
         fo.write(PNG_SIGNATURE)
-        for chunk in out_chunks:
+        for chunk in rebuilt:
             fo.write(chunk)
 
+    # sanity - size warning
+    if os.path.getsize(out_path) > 256:
+        print("⚠️  Output > 256 bytes; Kovaaks may still load it but smaller is safer.")
+
+    # optional validation / resize after writing (pixel data unchanged)
+    _validate_or_resize(out_path, auto_fix=False)  # ensure still 50×50 RGBA
+
     return out_path
 
 
-# ────────────────────────────────────────────────────────── gradio ui ──
-def process(src_file, tgt_file):
-    """Gradio wrapper → returns path to the merged PNG."""
+# ────────────────────────────────────────────────── Gradio UI ──
+def process(src_file: Path, tgt_file: Path, auto_fix: bool):
     if not src_file or not tgt_file:
-        raise gr.Error("Please upload both images.")
+        raise gr.Error("Upload both images first.")
+
+    # optional in-place auto-fix for target
+    try:
+        _validate_or_resize(str(tgt_file), auto_fix)
+    except Exception as e:
+        raise gr.Error(str(e))
 
     tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-    tmp.close()  # we only need the path
+    tmp.close()
     try:
-        clone_metadata(src_file, tgt_file, tmp.name)
+        clone_metadata(str(src_file), str(tgt_file), tmp.name, auto_fix=False)
     except Exception as e:
         os.remove(tmp.name)
         raise gr.Error(str(e))
@@ -84,33 +150,32 @@ def process(src_file, tgt_file):
     return tmp.name
 
 
-with gr.Blocks(title="Kovaaks PNG Metadata Cloner") as demo:
+with gr.Blocks(title="Kovaaks PNG Metadata Cloner v2") as demo:
     gr.Markdown(
         """
-        ## Kovaaks PNG Metadata Cloner 🏹
-
-        **Step 1.** Upload a *metadata source* PNG that already works in-game  
-        **Step 2.** Upload your *target* PNG (new 50 × 50 crosshair)  
-        **Step 3.** Click **Clone metadata** and download the result.
+        # 🏹 Kovaaks PNG Metadata Cloner v2
+        1. **Image 1 (source)** – a PNG that already works in-game  
+        2. **Image 2 (target)** – your new crosshair (any size/format)  
+        3. Tick **auto-fix** to force 50 × 50 RGBA conversion  
+        4. Click **Clone** → download the ready file
 
-        Put the output file into  
+        Put the result into  
         `%LOCALAPPDATA%\\FPSAimTrainer\\Saved\\MyImport\\crosshairs`
         """
     )
-
     with gr.Row():
-        src = gr.File(label="Image 1 – source (metadata)", type="filepath")
-        tgt = gr.File(label="Image 2 – target (pixels)",   type="filepath")
-
+        src = gr.File(label="Image 1 — metadata source (PNG)", type="filepath")
+        tgt = gr.File(label="Image 2 — target crosshair", type="filepath")
+    auto = gr.Checkbox(True, label="Auto-fix target to 50×50 RGBA")
     out = gr.File(label="Output PNG", interactive=False)
     btn = gr.Button("Clone metadata →")
 
-    btn.click(fn=process, inputs=[src, tgt], outputs=[out])
+    btn.click(fn=process, inputs=[src, tgt, auto], outputs=[out])
 
     gr.Markdown(
-        "✔️ Ancillary chunks (`sRGB`, `gAMA`, `pHYs`, `eXIf`, `tEXt`, …) are copied verbatim. "
-        "Pixel data from the target image is left untouched."
+        "Ancillary chunks (sRGB, gAMA, pHYs, eXIf, etc.) are copied verbatim, "
+        "duplicates removed, order normalised. Pixel data comes solely from Image 2."
     )
 
 if __name__ == "__main__":
-    demo.launch(show_api=False, share=False)
+    demo.launch(show_api=False)