S

Client/Scripts/hunt.py

@@ -207,15 +207,17 @@ def detect_exif_codes_from_files(
             ): post_id
             for post_id in post_ids
         }
-        for future in as_completed(futures):
-            if stop_event.is_set():
-                break
-            post_id = futures[future]
-            try:
-                code = future.result()
-            except Exception:
-                code = None
-            results[post_id] = code
+        with tqdm(total=len(futures), unit="image", desc="Scanning exif") as pbar:
+            for future in as_completed(futures):
+                if stop_event.is_set():
+                    break
+                post_id = futures[future]
+                try:
+                    code = future.result()
+                except Exception:
+                    code = None
+                results[post_id] = code
+                pbar.update(1)
     return results
 
 def chunked(seq: Sequence[str], size: int) -> Iterator[List[str]]:
@@ -520,7 +522,7 @@ try:
         if filtered:
             if DRY_RUN:
                 print("Dry run outputs (post_id -> page):")
-            scan_with_retries(filtered, phpsessid, conn, post_ids_dict, exif_types, "Scanning", stop_event)
+            scan_with_retries(filtered, phpsessid, conn, post_ids_dict, exif_types, "Scanning exif", stop_event)
         if DRY_RUN:
             continue
         exif_pending = {

Client/Scripts/scan_existing_exif.py

@@ -0,0 +1,217 @@
+import os
+import sqlite3
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+IMAGES_DIR = os.path.join(ROOT_DIR, "images")
+STASH_DIR = os.path.join(IMAGES_DIR, "Stash")
+DB_PATH = os.path.join(ROOT_DIR, "db.sqlite")
+MAX_WORKERS = min(16, os.cpu_count() or 8)
+EXIF_METADATA_MAX_BYTES = 512
+EXIF_TYPE_ORDER = ("novelai", "sd", "comfy", "mj", "celsys", "photoshop", "stealth")
+EXIF_TYPE_TO_CODE = {name: idx + 1 for idx, name in enumerate(EXIF_TYPE_ORDER)}
+PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+
+def open_db(path: str) -> sqlite3.Connection:
+    conn = sqlite3.connect(path)
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS pixif_cache (
+            post_id TEXT PRIMARY KEY,
+            url TEXT,
+            exif_type INTEGER
+        )
+        """
+    )
+    conn.commit()
+    ensure_db_schema(conn)
+    return conn
+
+def ensure_db_schema(conn: sqlite3.Connection) -> None:
+    columns = [row[1] for row in conn.execute("PRAGMA table_info(pixif_cache)")]
+    if "exif_type" not in columns:
+        conn.execute("ALTER TABLE pixif_cache ADD COLUMN exif_type INTEGER")
+        conn.commit()
+
+def determine_exif_type(metadata: Optional[bytes]) -> Optional[str]:
+    if metadata is None:
+        return None
+    if metadata == b"TitleAI generated image":
+        return "novelai"
+    if metadata.startswith(b"parameter"):
+        return "sd"
+    if b'{"' in metadata:
+        return "comfy"
+    if metadata.startswith(b"SoftwareCelsys"):
+        return "celsys"
+    return "photoshop"
+
+def exif_type_to_code(exif_type: Optional[str]) -> Optional[int]:
+    if not exif_type:
+        return None
+    return EXIF_TYPE_TO_CODE.get(exif_type)
+
+def parse_png_metadata(data: bytes) -> Optional[bytes]:
+    index = 8
+    while index < len(data):
+        if index + 8 > len(data):
+            break
+        chunk_len = int.from_bytes(data[index:index + 4], "big")
+        chunk_type = data[index + 4:index + 8]
+        index += 8
+        if chunk_type == b"tEXt":
+            content = data[index:index + chunk_len]
+            return content.replace(b"\0", b"")
+        if chunk_type == b"iTXt":
+            content = data[index:index + chunk_len]
+            return content.strip()
+        index += chunk_len + 4
+    return None
+
+def parse_png_metadata_file(path: str) -> Optional[bytes]:
+    try:
+        with open(path, "rb") as handle:
+            head = handle.read(EXIF_METADATA_MAX_BYTES)
+            if not head.startswith(PNG_SIGNATURE):
+                return None
+            return parse_png_metadata(head)
+    except Exception:
+        return None
+
+def byteize(alpha: np.ndarray) -> np.ndarray:
+    alpha = alpha.T.reshape((-1,))
+    alpha = alpha[:(alpha.shape[0] // 8) * 8]
+    alpha = np.bitwise_and(alpha, 1)
+    alpha = alpha.reshape((-1, 8))
+    alpha = np.packbits(alpha, axis=1)
+    return alpha
+
+class LSBExtractor:
+    def __init__(self, alpha: np.ndarray) -> None:
+        self.data = byteize(alpha)
+        self.pos = 0
+
+    def get_next_n_bytes(self, n: int) -> bytearray:
+        n_bytes = self.data[self.pos:self.pos + n]
+        self.pos += n
+        return bytearray(n_bytes)
+
+    def read_32bit_integer(self) -> Optional[int]:
+        bytes_list = self.get_next_n_bytes(4)
+        if len(bytes_list) == 4:
+            return int.from_bytes(bytes_list, byteorder="big")
+        return None
+
+def extract_stealth_metadata(image: Image.Image) -> bool:
+    if "A" not in image.getbands():
+        raise AssertionError("image format")
+    alpha = np.array(image.getchannel("A"))
+    reader = LSBExtractor(alpha)
+    magic = "stealth_pngcomp"
+    read_magic = reader.get_next_n_bytes(len(magic)).decode("utf-8")
+    if magic != read_magic:
+        raise AssertionError("magic number")
+    read_len = reader.read_32bit_integer()
+    if read_len is None:
+        raise AssertionError("length missing")
+    return True
+
+def has_stealth_png_path(path: str) -> bool:
+    try:
+        with Image.open(path) as image:
+            return extract_stealth_metadata(image)
+    except Exception:
+        return False
+
+def detect_exif_code_from_path(path: str) -> Optional[int]:
+    metadata = parse_png_metadata_file(path)
+    exif_type = determine_exif_type(metadata)
+    code = exif_type_to_code(exif_type)
+    if code is not None:
+        return code
+    if has_stealth_png_path(path):
+        return EXIF_TYPE_TO_CODE.get("stealth")
+    return None
+
+def fetch_pending_post_ids(conn: sqlite3.Connection) -> List[str]:
+    rows = conn.execute(
+        """
+        SELECT post_id
+        FROM pixif_cache
+        WHERE exif_type IS NULL
+          AND COALESCE(url, '') != ''
+        """
+    ).fetchall()
+    return [str(row[0]) for row in rows]
+
+def update_exif_types(conn: sqlite3.Connection, rows: Sequence[Tuple[int, str]]) -> None:
+    if not rows:
+        return
+    conn.executemany(
+        """
+        UPDATE pixif_cache SET exif_type = ?
+        WHERE post_id = ?
+        """,
+        rows,
+    )
+
+def detect_exif_codes_from_files(
+    post_ids: Sequence[str],
+    stash_dir: str,
+    max_workers: int = MAX_WORKERS,
+) -> Dict[str, Optional[int]]:
+    if not post_ids:
+        return {}
+    results: Dict[str, Optional[int]] = {}
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = {
+            executor.submit(
+                detect_exif_code_from_path,
+                os.path.join(stash_dir, f"{post_id}.png"),
+            ): post_id
+            for post_id in post_ids
+        }
+        with tqdm(total=len(futures), unit="image", desc="Scanning exif") as pbar:
+            for future in as_completed(futures):
+                post_id = futures[future]
+                try:
+                    code = future.result()
+                except Exception:
+                    code = None
+                results[post_id] = code
+                pbar.update(1)
+    return results
+
+def main() -> int:
+    os.makedirs(STASH_DIR, exist_ok=True)
+    conn = open_db(DB_PATH)
+    try:
+        post_ids = fetch_pending_post_ids(conn)
+        if not post_ids:
+            print("No pending rows.")
+            return 0
+        existing = [post_id for post_id in post_ids if os.path.exists(os.path.join(STASH_DIR, f"{post_id}.png"))]
+        if not existing:
+            print("No matching images in stash.")
+            return 0
+        results = detect_exif_codes_from_files(existing, STASH_DIR)
+        rows = [
+            (exif_type, post_id)
+            for post_id, exif_type in results.items()
+            if exif_type is not None
+        ]
+        if rows:
+            with conn:
+                update_exif_types(conn, rows)
+        print(f"Updated {len(rows)} rows.")
+        return 0
+    finally:
+        conn.close()
+
+if __name__ == "__main__":
+    raise SystemExit(main())

Client/main.py

@@ -32,14 +32,15 @@ MENU_SECTIONS: Sequence[Tuple[str, Sequence[Tuple[str, str, str]]]] = (
         "Hunt",
         (
             ("5", "Hunt", os.path.join(SCRIPTS_DIR, "hunt.py")),
+            ("6", "Scan existing exif", os.path.join(SCRIPTS_DIR, "scan_existing_exif.py")),
         ),
     ),
     (
         "Maintenance",
         (
-            ("6", "Clear text logs", os.path.join(SCRIPTS_DIR, "clear_texts.py")),
-            ("7", "Clear DB empty rows", os.path.join(SCRIPTS_DIR, "clear_db.py")),
-            ("8", "Clear images (except Stash)", os.path.join(SCRIPTS_DIR, "clear_images.py")),
+            ("7", "Clear text logs", os.path.join(SCRIPTS_DIR, "clear_texts.py")),
+            ("8", "Clear DB empty rows", os.path.join(SCRIPTS_DIR, "clear_db.py")),
+            ("9", "Clear images (except Stash)", os.path.join(SCRIPTS_DIR, "clear_images.py")),
         ),
     ),
 )