142 lines
4.9 KiB
Python
142 lines
4.9 KiB
Python
"""photo-pipeline: photo-ingest flow (v3 — single-pass, checkpointed).
|
|
|
|
Hashes incoming images, checks the persistent fingerprint DB (exact dups via
|
|
sha256, multi-path dup tracking via known_paths), and registers new files.
|
|
|
|
Design: ONE walk of the tree. Each file: path-known → skip fast; sha-known at
|
|
another path → record path in known_paths (dup); else register new. Inserts
|
|
commit every `batch_size` rows (WAL) = checkpoint. Single pass = O(n), no
|
|
re-walk, no O(n²), no engine busy-loop.
|
|
|
|
Performance: dhash-only rows (~6ms/file); phash only when check_near_dups.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
import signal
|
|
import time as _time
|
|
|
|
from prefect import flow, task
|
|
from PIL import Image
|
|
|
|
import photo_db as db
|
|
|
|
EXT_IMAGES = {".jpg", ".jpeg", ".png", ".heic", ".webp", ".gif", ".tif", ".tiff", ".bmp"}
|
|
|
|
|
|
def _is_image(p: Path) -> bool:
|
|
return p.is_file() and p.suffix.lower() in EXT_IMAGES
|
|
|
|
|
|
def walk_files(base_dir: str):
|
|
"""Stream image files under base_dir (generator — memory-safe for 50K+)."""
|
|
root = Path(base_dir)
|
|
if not root.exists():
|
|
raise FileNotFoundError(f"{root} does not exist")
|
|
for p in root.rglob("*"):
|
|
if _is_image(p):
|
|
yield str(p)
|
|
|
|
|
|
@task
|
|
def scan_once(base_dir: str, source: str, batch_size: int, check_near_dups: bool) -> dict:
|
|
"""One walk of the tree; classify + register in a single pass."""
|
|
import time
|
|
|
|
db.init_db()
|
|
conn = db.get_db()
|
|
conn.execute(
|
|
"CREATE TABLE IF NOT EXISTS known_paths (path TEXT PRIMARY KEY, sha256 TEXT NOT NULL)"
|
|
)
|
|
totals = {"exact_dups": 0, "near_dups": 0, "new": 0, "known": 0, "skipped": 0}
|
|
pending = []
|
|
t0 = time.time()
|
|
|
|
def _timeout_handler(signum, frame):
|
|
raise TimeoutError(f"file took >{FILE_TIMEOUT}s")
|
|
|
|
FILE_TIMEOUT = 10 # seconds per file — guards against hung/corrupt files
|
|
signal.signal(signal.SIGALRM, _timeout_handler)
|
|
|
|
for p_str in walk_files(base_dir):
|
|
p = Path(p_str)
|
|
try:
|
|
signal.alarm(FILE_TIMEOUT)
|
|
known = conn.execute(
|
|
"SELECT 1 FROM image_hashes WHERE path=? UNION SELECT 1 FROM known_paths WHERE path=?",
|
|
(p_str, p_str),
|
|
).fetchone()
|
|
if known:
|
|
totals["known"] += 1
|
|
signal.alarm(0)
|
|
continue
|
|
|
|
sha = db.sha256_file(p_str)
|
|
match = conn.execute(
|
|
"SELECT path FROM image_hashes WHERE sha256=?", (sha,)
|
|
).fetchone()
|
|
if match:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO known_paths (path, sha256) VALUES (?,?)", (p_str, sha))
|
|
totals["exact_dups"] += 1
|
|
signal.alarm(0)
|
|
continue
|
|
|
|
hx = db.hash_image(p)
|
|
with Image.open(p) as im:
|
|
w, h = im.size
|
|
pending.append((sha, hx["phash"], hx["dhash"], p.stat().st_size,
|
|
w, h, str(p), source))
|
|
totals["new"] += 1
|
|
signal.alarm(0)
|
|
|
|
if len(pending) >= batch_size:
|
|
conn.executemany(
|
|
"INSERT INTO image_hashes (sha256, phash, dhash, file_size, width, height, path, source) "
|
|
"VALUES (?,?,?,?,?,?,?,?)", pending)
|
|
conn.commit()
|
|
print(f" checkpoint: {totals['new']} new ({_time.time()-t0:.0f}s)", flush=True)
|
|
pending = []
|
|
|
|
except TimeoutError as e:
|
|
totals["skipped"] += 1
|
|
print(f" TIMEOUT {p.name}: {e}", flush=True)
|
|
except Exception as e:
|
|
totals["skipped"] += 1
|
|
print(f" SKIP {p.name}: {type(e).__name__}: {e}", flush=True)
|
|
finally:
|
|
signal.alarm(0)
|
|
|
|
if pending:
|
|
conn.executemany(
|
|
"INSERT INTO image_hashes (sha256, phash, dhash, file_size, width, height, path, source) "
|
|
"VALUES (?,?,?,?,?,?,?,?)", pending)
|
|
conn.commit()
|
|
conn.close()
|
|
print(f"scan_once done in {time.time()-t0:.0f}s: {totals}", flush=True)
|
|
return totals
|
|
|
|
|
|
@flow(name="photo-ingest")
|
|
def photo_ingest(base_dir: str, source: str = "takeout", batch_size: int = 1000,
|
|
max_batches: int | None = None, check_near_dups: bool = False):
|
|
"""Single-pass hash + dedup scan of a folder. One walk, checkpointed inserts.
|
|
|
|
Args:
|
|
base_dir: folder to scan
|
|
source: label (e.g. takeout, archive-pictures)
|
|
batch_size: insert checkpoint interval (default 1000)
|
|
max_batches: kept for compatibility (unused — single pass)
|
|
check_near_dups: reserved (dhash-only rows; near-dup off)
|
|
"""
|
|
return scan_once(base_dir, source, batch_size, check_near_dups)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
d = sys.argv[1] if len(sys.argv) > 1 else "/tmp/sample"
|
|
src = sys.argv[2] if len(sys.argv) > 2 else "cli-test"
|
|
bs = int(sys.argv[3]) if len(sys.argv) > 3 else 1000
|
|
r = photo_ingest(d, source=src, batch_size=bs)
|
|
print(r)
|