116 lines
3.2 KiB
Python
116 lines
3.2 KiB
Python
"""photo-pipeline: photo-ingest flow (v1).
|
|
|
|
Stage 1 of the pipeline: hash incoming images, check against the persistent
|
|
fingerprint DB (exact + near dupes), and register new ones.
|
|
|
|
Run via Prefect deployment on photo-pool (see prefect.yaml).
|
|
"""
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
from prefect import flow, task
|
|
|
|
import photo_db as db
|
|
|
|
|
|
@task
|
|
def scan_directory(base_dir: str) -> list[str]:
|
|
"""Enumerate image files in a directory tree."""
|
|
exts = {".jpg", ".jpeg", ".png", ".heic", ".webp", ".gif", ".tif", ".tiff", ".bmp"}
|
|
root = Path(base_dir)
|
|
if not root.exists():
|
|
raise FileNotFoundError(f"{root} does not exist")
|
|
found = [
|
|
str(p)
|
|
for p in root.rglob("*")
|
|
if p.is_file() and p.suffix.lower() in exts
|
|
]
|
|
print(f"Found {len(found)} images under {root}")
|
|
return found
|
|
|
|
|
|
@task
|
|
def check_duplicates(image_paths: list[str]) -> dict:
|
|
"""Check each image against the fingerprint DB. Returns classification."""
|
|
db.init_db()
|
|
conn = db.get_db()
|
|
exact_dups = []
|
|
near_dups = []
|
|
new_images = []
|
|
for p_str in image_paths:
|
|
p = Path(p_str)
|
|
try:
|
|
match = db.check_exact_dup(conn, p)
|
|
if match:
|
|
exact_dups.append((p_str, match))
|
|
continue
|
|
near, dist = db.check_near_dup(conn, p)
|
|
if near:
|
|
near_dups.append((p_str, near, dist))
|
|
continue
|
|
new_images.append(p_str)
|
|
except Exception as e:
|
|
print(f" SKIP {p.name}: {e}")
|
|
conn.close()
|
|
result = {
|
|
"total": len(image_paths),
|
|
"exact_dups": len(exact_dups),
|
|
"near_dups": len(near_dups),
|
|
"new": len(new_images),
|
|
"exact_dup_list": exact_dups[:50],
|
|
"near_dup_list": near_dups[:50],
|
|
"new_list": new_images,
|
|
}
|
|
print(
|
|
f"Check: {result['total']} total, "
|
|
f"{result['exact_dups']} exact dups, {result['near_dups']} near dups, "
|
|
f"{result['new']} new"
|
|
)
|
|
return result
|
|
|
|
|
|
@task
|
|
def register_new_images(image_paths: list[str], source: str) -> int:
|
|
"""Add hashes for confirmed-new images into the fingerprint DB."""
|
|
db.init_db()
|
|
conn = db.get_db()
|
|
registered = 0
|
|
for p_str in image_paths:
|
|
p = Path(p_str)
|
|
try:
|
|
if db.register_image(conn, p, source=source):
|
|
registered += 1
|
|
except Exception as e:
|
|
print(f" FAIL register {p.name}: {e}")
|
|
conn.commit()
|
|
conn.close()
|
|
print(f"Registered {registered} new images (source={source})")
|
|
return registered
|
|
|
|
|
|
@flow(name="photo-ingest")
|
|
def photo_ingest(base_dir: str, source: str = "takeout", register: bool = True):
|
|
"""Hash + dedup-check a folder against the persistent library."""
|
|
images = scan_directory(base_dir)
|
|
if not images:
|
|
print("No images found — nothing to do.")
|
|
return {"total": 0}
|
|
|
|
result = check_duplicates(images)
|
|
|
|
if register and result["new_list"]:
|
|
n = register_new_images(result["new_list"], source=source)
|
|
result["registered"] = n
|
|
|
|
return result
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Local run (no deployment)
|
|
import sys
|
|
|
|
d = sys.argv[1] if len(sys.argv) > 1 else "/tmp/sample"
|
|
r = photo_ingest(d)
|
|
print(r)
|