ray
117 lines
4.8 KiB
Python
117 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Immich YOLO Photo Classifier Pipeline
|
|
Classifies photos as: has-people, scenic, or neither.
|
|
Moves non-people/non-scenic photos to NO PEOPLE PHOTOS folder.
|
|
|
|
Usage:
|
|
/mnt/storage/yolo_venv/bin/python3 this_script.py
|
|
|
|
Requirements:
|
|
- /mnt/storage/yolo_venv/ with ultralytics installed
|
|
- ~/yolov8n_openvino_model/ (exported from yolov8n.pt)
|
|
- Immich PostgreSQL container running
|
|
- Photos on /mnt/wd-passport/immich/photos/
|
|
"""
|
|
import multiprocessing, os, time, warnings, subprocess, shutil
|
|
warnings.filterwarnings("ignore")
|
|
|
|
PHOTO_BASE = "/mnt/wd-passport/immich/photos"
|
|
NO_PEOPLE_DIR = "/mnt/wd-passport/immich/NO PEOPLE PHOTOS"
|
|
DB_QUERY = '''docker exec immich_postgres psql -U postgres -d immich -t -A -c "SELECT \\\"originalPath\\\" FROM asset WHERE type='IMAGE' AND visibility='timeline' AND \\\"originalPath\\\" LIKE '/data/library/%' ORDER BY \\\"originalPath\\\";" 2>/dev/null'''
|
|
|
|
PERSON_CLASS = 0
|
|
NATURE_CLASSES = {16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 58, 77, 80}
|
|
MOVE_CLASSES = {56, 57, 59, 60, 61, 62, 63, 64, 65, 67, 70, 71, 72, 73, 74, 75, 76, 78, 79,
|
|
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
|
|
|
|
def worker(args):
|
|
model_path, paths = args
|
|
from ultralytics import YOLO
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
model = YOLO(model_path, task="detect")
|
|
results = []
|
|
for ap in paths:
|
|
try:
|
|
r = model(ap, verbose=False)
|
|
has_person = has_nature = has_move = False
|
|
for r2 in r:
|
|
if r2.boxes:
|
|
for box in r2.boxes:
|
|
cls = int(box.cls[0])
|
|
if cls == PERSON_CLASS: has_person = True
|
|
elif cls in NATURE_CLASSES: has_nature = True
|
|
elif cls in MOVE_CLASSES: has_move = True
|
|
if has_person: results.append((ap, "keep_people"))
|
|
elif has_nature: results.append((ap, "keep_scenic"))
|
|
elif has_move: results.append((ap, "move"))
|
|
else: results.append((ap, "check_size"))
|
|
except Exception as e:
|
|
results.append((ap, f"error:{e}"))
|
|
return results
|
|
|
|
if __name__ == "__main__":
|
|
multiprocessing.set_start_method("forkserver", force=True)
|
|
num_workers = min(multiprocessing.cpu_count(), 8)
|
|
os.makedirs(NO_PEOPLE_DIR, exist_ok=True)
|
|
model_path = os.path.expanduser("~/yolov8n_openvino_model/")
|
|
|
|
print(f"[1] Querying DB...", flush=True)
|
|
result = subprocess.run(DB_QUERY, shell=True, capture_output=True, text=True, timeout=120)
|
|
db_paths = [p.strip() for p in result.stdout.strip().split('\n') if p.strip()]
|
|
|
|
print(f"[2] Mapping {len(db_paths)} paths...", flush=True)
|
|
actual = []
|
|
for p in db_paths:
|
|
ap = PHOTO_BASE + "/" + p[6:]
|
|
if os.path.exists(ap):
|
|
actual.append(ap)
|
|
print(f" {len(actual)} accessible files", flush=True)
|
|
|
|
print(f"[3] Classifying with {num_workers} workers...", flush=True)
|
|
chunks = [actual[i::num_workers] for i in range(num_workers)]
|
|
chunk_args = [(model_path, c) for c in chunks]
|
|
t_start = time.time()
|
|
with multiprocessing.Pool(num_workers) as pool:
|
|
all_results = pool.map(worker, chunk_args)
|
|
elapsed = time.time() - t_start
|
|
|
|
keep_people = keep_scenic = errors = 0
|
|
to_move, check_size = [], []
|
|
for chunk in all_results:
|
|
for ap, d in chunk:
|
|
if d == "keep_people": keep_people += 1
|
|
elif d == "keep_scenic": keep_scenic += 1
|
|
elif d == "move": to_move.append(ap)
|
|
elif d == "check_size": check_size.append(ap)
|
|
else: errors += 1
|
|
|
|
print(f"\n YOLO: {elapsed/60:.1f}min ({len(actual)/elapsed:.1f} img/s)", flush=True)
|
|
print(f" People:{keep_people} Scenic:{keep_scenic} Move:{len(to_move)} Check:{len(check_size)} Err:{errors}", flush=True)
|
|
|
|
print(f"[4] Size check on {len(check_size)} undetected...", flush=True)
|
|
for ap in check_size:
|
|
try:
|
|
if os.path.getsize(ap) < 400000:
|
|
to_move.append(ap)
|
|
except: pass
|
|
print(f" Now {len(to_move)} to move", flush=True)
|
|
|
|
print(f"[5] Moving {len(to_move)} photos...", flush=True)
|
|
moved = file_errors = 0
|
|
for ap in to_move:
|
|
try:
|
|
rel = os.path.relpath(ap, PHOTO_BASE)
|
|
dest = os.path.join(NO_PEOPLE_DIR, rel)
|
|
os.makedirs(os.path.dirname(dest), exist_ok=True)
|
|
shutil.move(ap, dest)
|
|
moved += 1
|
|
except Exception as e:
|
|
file_errors += 1
|
|
|
|
print(f"\nDONE! Moved {moved} photos", flush=True)
|
|
print(f" Total:{len(actual)} Kept(people):{keep_people} Kept(scenic):{keep_scenic} Moved:{moved} Errors:{file_errors}", flush=True)
|
|
print(f" Time: {(time.time()-t_start)/60:.1f}min", flush=True)
|
|
print(f" Output: {NO_PEOPLE_DIR}", flush=True)
|