create_cluster.py

import os
from PIL import Image
import torch
from torchvision import models, transforms
from sklearn.cluster import KMeans
import numpy as np
from tqdm import tqdm
from utils.utils import ensure_dir
import argparse


def main():
    argparser = argparse.ArgumentParser(
        description="Cluster images using ResNet50 features"
    )
    argparser.add_argument(
        "--data-dir",
        type=str,
        required=True,
        help="Path to the folder containing TRAINING images to cluster",
    )
    argparser.add_argument(
        "--output-dir",
        type=str,
        default="./",
        help="Output dir to save cluster assignments",
    )
    args = argparser.parse_args()

    # Path to your image folder
    image_folder = args.data_dir
    if not os.path.exists(image_folder):
        raise ValueError(f"Image folder '{image_folder}' does not exist.")

    ensure_dir(args.output_dir)

    # Preprocessing for ResNet
    preprocess = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )

    # Load pre-trained ResNet50
    device = torch.device(
        "mps"
        if torch.backends.mps.is_available()
        else ("cuda" if torch.cuda.is_available() else "cpu")
    )
    print("using device:", device)
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
    model = torch.nn.Sequential(
        *(list(model.children())[:-1])
    )  # Remove final classification layer
    model = model.to(device)
    model.eval()

    # Extract features with progress bar
    features = []
    image_paths = []
    with torch.no_grad():
        for fname in tqdm(os.listdir(image_folder), desc="Extracting features"):
            if fname.lower().endswith((".png", ".jpg", ".jpeg")):
                img_path = os.path.join(image_folder, fname)
                img = Image.open(img_path).convert("RGB")
                input_tensor = preprocess(img).unsqueeze(0).to(device)
                output = model(input_tensor)
                features.append(output.cpu().squeeze().numpy())
                image_paths.append(img_path)

    if len(features) == 0:
        print(
            "No images found! Did you specify the train folder correctly? It must point to the train folder!"
        )
        exit()

    features = np.array(features).reshape(len(features), -1)

    print("Start clustering...")
    kmeans = KMeans(n_clusters=360, random_state=42, n_init="auto")
    kmeans.fit(features)

    # save cluster assignments
    cluster_assignments = kmeans.labels_
    output_file = os.path.join(args.output_dir, "cluster.txt")
    with open(output_file, "w") as f:
        for img_path, cluster in zip(image_paths, cluster_assignments):
            f.write(f"{cluster}/{img_path.split('/')[-1]}\n")


if __name__ == "__main__":
    main()