Rem from hunt.py

Client/Scripts/hunt.py

@@ -7,10 +7,8 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import shutil
 from typing import Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
-import numpy as np
 import requests
 from tqdm import tqdm
-from PIL import Image
 
 LOCAL = 0
 DRY_RUN = 0
@@ -32,11 +30,6 @@ STREAM_IDLE_TIMEOUT_SECONDS = 45
 STREAM_MAX_READ_TIMEOUTS = 3
 STREAM_MAX_RETRIES = 3
 STREAM_RETRY_DELAY_SECONDS = 2
-EXIF_BATCH_SIZE = 200
-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)}
-EXIF_METADATA_MAX_BYTES = 512
-PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
 stop_event = threading.Event()
 
 def read_dotenv_value(path: str, key: str) -> Optional[str]:
@@ -75,153 +68,13 @@ def open_db(path: str) -> sqlite3.Connection:
         """
         CREATE TABLE IF NOT EXISTS pixif_cache (
             post_id TEXT PRIMARY KEY,
-            url TEXT,
-            exif_type INTEGER
+            url TEXT
         )
         """
     )
     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 b"Dig" in metadata:
-        return "mj"
-    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 detect_exif_codes_from_files(
-    post_ids: Sequence[str],
-    stash_dir: str,
-    max_workers: int = MAX_WORKERS,
-    stop_event: threading.Event = stop_event,
-) -> 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):
-                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]]:
     for i in range(0, len(seq), size):
         yield seq[i:i + size]
@@ -229,20 +82,18 @@ def chunked(seq: Sequence[str], size: int) -> Iterator[List[str]]:
 def fetch_cached_state(
     conn: sqlite3.Connection,
     post_ids: Sequence[str],
-) -> Tuple[Dict[str, Optional[str]], Dict[str, Optional[int]]]:
+) -> Dict[str, Optional[str]]:
     post_ids_dict = {post_id: None for post_id in post_ids}
-    exif_types = {post_id: None for post_id in post_ids}
     if not post_ids:
-        return post_ids_dict, exif_types
+        return post_ids_dict
 
     for chunk in chunked(post_ids, 900):
         placeholders = ",".join("?" for _ in chunk)
-        query = f"SELECT post_id, COALESCE(url, ''), exif_type FROM pixif_cache WHERE post_id IN ({placeholders})"
-        for post_id, url, exif_type in conn.execute(query, chunk):
+        query = f"SELECT post_id, COALESCE(url, '') FROM pixif_cache WHERE post_id IN ({placeholders})"
+        for post_id, url in conn.execute(query, chunk):
             post_ids_dict[post_id] = url
-            exif_types[post_id] = exif_type
 
-    return post_ids_dict, exif_types
+    return post_ids_dict
 
 def upsert_urls(conn: sqlite3.Connection, rows: Sequence[Tuple[str, str]]) -> None:
     if not rows:
@@ -256,30 +107,11 @@ def upsert_urls(conn: sqlite3.Connection, rows: Sequence[Tuple[str, str]]) -> No
         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 fetch_exif_types(
-    post_urls: Dict[str, str],
-    phpsessid: str,
-    stop_event: threading.Event = stop_event,
-) -> Dict[str, Optional[int]]:
-    return {}
-
 def stream_pixif_updates(
     post_ids: Sequence[str],
     phpsessid: str,
     conn: sqlite3.Connection,
     post_ids_dict: Dict[str, Optional[str]],
-    exif_types: Dict[str, Optional[int]],
     desc: str,
     stop_event: threading.Event = stop_event,
 ) -> int:
@@ -330,7 +162,6 @@ def stream_pixif_updates(
                                 with conn:
                                     upsert_urls(conn, [(post_id, url)])
                             post_ids_dict[post_id] = url
-                            exif_types[post_id] = exif_types.get(post_id)
                             pbar.update(1)
                             processed += 1
                             last_progress = now
@@ -464,7 +295,6 @@ def scan_with_retries(
     phpsessid: str,
     conn: sqlite3.Connection,
     post_ids_dict: Dict[str, Optional[str]],
-    exif_types: Dict[str, Optional[int]],
     desc: str,
     stop_event: threading.Event = stop_event,
 ) -> None:
@@ -473,7 +303,7 @@ def scan_with_retries(
     remaining = list(post_ids)
     attempts = 0
     while remaining and attempts < STREAM_MAX_RETRIES and not stop_event.is_set():
-        stream_pixif_updates(remaining, phpsessid, conn, post_ids_dict, exif_types, desc, stop_event)
+        stream_pixif_updates(remaining, phpsessid, conn, post_ids_dict, desc, stop_event)
         remaining = [post_id for post_id in remaining if post_ids_dict.get(post_id) is None]
         if not remaining:
             break
@@ -489,7 +319,7 @@ try:
         with open(os.path.join(ROOT_DIR, valid[index]), "r") as f:
             post_ids = f.read().split()
 
-        post_ids_dict, exif_types = fetch_cached_state(conn, post_ids)
+        post_ids_dict = fetch_cached_state(conn, post_ids)
         if DRY_RUN:
             filtered = list(post_ids)
         else:
@@ -503,26 +333,9 @@ 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 exif", stop_event)
+            scan_with_retries(filtered, phpsessid, conn, post_ids_dict, "API hunt", stop_event)
         if DRY_RUN:
             continue
-        exif_pending = {
-            post_id: decode_if_binary(url)
-            for post_id, url in post_ids_dict.items()
-            if url and exif_types.get(post_id) is None
-        }
-        if exif_pending:
-            exif_results = fetch_exif_types(exif_pending, phpsessid, stop_event=stop_event)
-            rows = [
-                (exif_type, post_id)
-                for post_id, exif_type in exif_results.items()
-                if exif_type is not None
-            ]
-            if rows:
-                with conn:
-                    update_exif_types(conn, rows)
-            for post_id, exif_type in exif_results.items():
-                exif_types[post_id] = exif_type
         to_download = {post_id: decode_if_binary(url) for post_id, url in post_ids_dict.items() if url and f"{post_id}.png" not in images_cache}
 
         if to_download:
@@ -532,24 +345,6 @@ try:
 
         images_cache.update(os.listdir(STASH_DIR))
 
-        local_pending = [
-            post_id
-            for post_id in post_ids
-            if exif_types.get(post_id) is None and f"{post_id}.png" in images_cache
-        ]
-        if local_pending:
-            local_results = detect_exif_codes_from_files(local_pending, STASH_DIR, stop_event=stop_event)
-            rows = [
-                (exif_type, post_id)
-                for post_id, exif_type in local_results.items()
-                if exif_type is not None
-            ]
-            if rows:
-                with conn:
-                    update_exif_types(conn, rows)
-            for post_id, exif_type in local_results.items():
-                exif_types[post_id] = exif_type
-
         print("Linking images to the group directory...")
         for i, post_id in enumerate(post_ids):
             stash_path = os.path.join(STASH_DIR, f"{post_id}.png")