chore: Update everything

cluster_images.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env -S uv run
+# /// script
+# dependencies = [
+#   "torch",
+#   "torchvision",
+#   "timm",
+#   "scikit-learn",
+#   "pillow",
+#   "imagehash",
+# ]
+# ///
+"""Group similar large images using ResNet embeddings or perceptual hashing."""
+
+import torch
+import timm
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from sklearn.cluster import KMeans, DBSCAN
+from sklearn.preprocessing import StandardScaler
+import imagehash
+import argparse
+
+# Load pretrained ResNet50 on M1
+device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+model = timm.create_model("resnet50", pretrained=True, num_classes=0)
+model = model.to(device)
+model.eval()
+
+
+def get_embedding(img_path, size=224):
+    """Extract embedding from image."""
+    try:
+        img = Image.open(img_path).convert("RGB")
+        img.thumbnail((size, size), Image.Resampling.LANCZOS)
+
+        canvas = Image.new("RGB", (size, size), (128, 128, 128))
+        offset = ((size - img.width) // 2, (size - img.height) // 2)
+        canvas.paste(img, offset)
+
+        x = torch.tensor(np.array(canvas), dtype=torch.float32)
+        x = x.permute(2, 0, 1) / 255.0
+        x = (x - torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)) / torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+
+        with torch.no_grad():
+            embedding = model(x.unsqueeze(0).to(device)).squeeze().cpu().numpy()
+        return embedding
+    except Exception as e:
+        print(f"Error processing {img_path}: {e}")
+        return None
+
+
+def get_phash(img_path):
+    """Extract perceptual hash from image."""
+    try:
+        img = Image.open(img_path).convert("RGB")
+        return imagehash.phash(img)
+    except Exception as e:
+        print(f"Error processing {img_path}: {e}")
+        return None
+
+
+def hash_distance(h1, h2):
+    """Hamming distance between two hashes."""
+    return h1 - h2
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Group similar images using embeddings or perceptual hashing.")
+    parser.add_argument("paths", nargs="+", type=Path, help="Image file(s) or directory")
+    parser.add_argument("--cluster", action="store_true", help="Use ResNet embedding clustering")
+    parser.add_argument("--perceptual-hash", action="store_true", help="Use perceptual hashing")
+    parser.add_argument("--clusters", type=int, help="Number of clusters for embedding mode (auto if not specified)")
+    parser.add_argument("--hash-threshold", type=int, default=5, help="Hamming distance threshold for perceptual hash")
+    args = parser.parse_args()
+
+    if not args.cluster and not args.perceptual_hash:
+        parser.error("Either --cluster or --perceptual-hash must be specified")
+    if args.cluster and args.perceptual_hash:
+        parser.error("Cannot specify both --cluster and --perceptual-hash")
+
+    return args
+
+
+def group_by_hash(valid_files, hashes, threshold):
+    """Group images by perceptual hash similarity."""
+    labels = [-1] * len(valid_files)
+    cluster_id = 0
+
+    for i in range(len(valid_files)):
+        if labels[i] != -1:
+            continue
+        labels[i] = cluster_id
+        for j in range(i + 1, len(valid_files)):
+            if labels[j] == -1 and hash_distance(hashes[i], hashes[j]) <= threshold:
+                labels[j] = cluster_id
+        cluster_id += 1
+
+    return np.array(labels)
+
+
+def main():
+    args = parse_args()
+
+    # Collect image files
+    img_files = []
+    for path in args.paths:
+        if path.is_dir():
+            img_files.extend(path.glob("*.[jJ][pP][gG]"))
+            img_files.extend(path.glob("*.[pP][nN][gG]"))
+        else:
+            img_files.append(path)
+
+    if not img_files:
+        print("No images found.")
+        return
+
+    if args.perceptual_hash:
+        # Perceptual hash mode
+        hashes = []
+        valid_files = []
+
+        for i, img_path in enumerate(img_files):
+            print(f"Processing {i + 1}/{len(img_files)}: {img_path.name}")
+            h = get_phash(img_path)
+            if h is not None:
+                hashes.append(h)
+                valid_files.append(img_path)
+
+        labels = group_by_hash(valid_files, hashes, args.hash_threshold)
+        n_clusters = len(np.unique(labels))
+
+    else:
+        # Embedding clustering mode
+        embeddings = []
+        valid_files = []
+
+        for i, img_path in enumerate(img_files):
+            print(f"Processing {i + 1}/{len(img_files)}: {img_path.name}")
+            emb = get_embedding(img_path)
+            if emb is not None:
+                embeddings.append(emb)
+                valid_files.append(img_path)
+
+        embeddings = np.array(embeddings)
+        scaler = StandardScaler()
+        embeddings = scaler.fit_transform(embeddings)
+
+        n_clusters = args.clusters or max(2, int(np.sqrt(len(embeddings) / 2)))
+        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
+        labels = kmeans.fit_predict(embeddings)
+
+    # Save clustered images
+    for img_path, cluster_label in zip(valid_files, labels):
+        new_name = f"group{cluster_label}_{img_path.name}"
+        dest = Path.cwd() / new_name
+        img_path.rename(dest) if img_path.parent == Path.cwd() else dest.write_bytes(img_path.read_bytes())
+
+    print(f"Clustered {len(valid_files)} images into {n_clusters} clusters in {Path.cwd()}")
+
+
+if __name__ == "__main__":
+    main()