code/build_grouped_split.py

"""
Phase 1.1 — Group-aware split builder for Augmented dataset.

Strategy: perceptual hash (pHash) every image in both Original and Augmented datasets,
then group images by Hamming-distance < threshold. Result: each unique source image
(plus all of its augmented derivatives) gets one group_id. We then do a stratified
GroupKFold split so all derivatives of a source stay on the same side.
"""
import argparse, json, os, sys
from pathlib import Path
from collections import defaultdict
from PIL import Image
import imagehash
import numpy as np
from sklearn.model_selection import StratifiedGroupKFold, train_test_split
from tqdm import tqdm


# Maps original-dataset class folder name -> canonical class label.
# (Augmented dataset uses slightly different folder names for some classes.)
CLASS_CANON = {
    "Central Serous Chorioretinopathy [Color Fundus]": "CSC",
    "Diabetic Retinopathy": "DR",
    "Disc Edema": "DiscEdema",
    "Glaucoma": "Glaucoma",
    "Healthy": "Healthy",
    "Macular Scar": "MacularScar",
    "Myopia": "Myopia",
    "Pterygium": "Pterygium",
    "Retinal Detachment": "RetinalDet",
    "Retinitis Pigmentosa": "RetinitisPig",
}


def list_images(root: Path):
    """Yield (path, class_canon) for every image."""
    out = []
    for class_dir in sorted(root.iterdir()):
        if not class_dir.is_dir():
            continue
        canon = CLASS_CANON.get(class_dir.name, class_dir.name)
        for img in sorted(class_dir.iterdir()):
            if img.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp"}:
                out.append((str(img), canon))
    return out


def phash_image(path, hash_size=8):
    try:
        with Image.open(path) as im:
            im = im.convert("RGB")
            return imagehash.phash(im, hash_size=hash_size)
    except Exception as e:
        print(f"  hash error {path}: {e}", file=sys.stderr)
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--original-dir", default="Database/Original_Dataset")
    ap.add_argument("--augmented-dir", default="Database/Augmented_Dataset")
    ap.add_argument("--output", default="holdout_split_augmented.json")
    ap.add_argument("--hamming-threshold", type=int, default=8,
                    help="pHash Hamming distance for considering two images near-duplicates (8/64 bits)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--n-folds", type=int, default=5)
    ap.add_argument("--test-frac", type=float, default=0.15)
    ap.add_argument("--val-frac", type=float, default=0.15)
    args = ap.parse_args()

    orig_imgs = list_images(Path(args.original_dir))
    aug_imgs = list_images(Path(args.augmented_dir))
    print(f"original: {len(orig_imgs)} images")
    print(f"augmented: {len(aug_imgs)} images")

    # Compute hashes
    print("\nHashing original dataset ...")
    orig_hashes = []
    for p, c in tqdm(orig_imgs):
        h = phash_image(p)
        if h is not None:
            orig_hashes.append((p, c, h))

    print("\nHashing augmented dataset ...")
    aug_hashes = []
    for p, c in tqdm(aug_imgs):
        h = phash_image(p)
        if h is not None:
            aug_hashes.append((p, c, h))

    # Each original image becomes its own group (group_id = orig index)
    # Each augmented image is assigned to the nearest original IN THE SAME CLASS
    # (constrains search and avoids cross-class matches due to vignette).
    # If the nearest original is further than `hamming_threshold`, the augmented
    # image becomes its own standalone group.
    print(f"\nGrouping augmented images to originals (Hamming <= {args.hamming_threshold}) ...")

    orig_by_class = defaultdict(list)  # class -> list of (idx_in_global, path, hash)
    for i, (p, c, h) in enumerate(orig_hashes):
        orig_by_class[c].append((i, p, h))

    groups = {}             # path -> group_id
    group_class = {}        # group_id -> class
    next_standalone_id = len(orig_hashes)

    # Originals: trivially their own group
    for i, (p, c, _) in enumerate(orig_hashes):
        groups[p] = i
        group_class[i] = c

    # Augmented: nearest-original lookup within same class
    matched, standalone = 0, 0
    for p, c, h in tqdm(aug_hashes):
        cands = orig_by_class.get(c, [])
        if not cands:
            groups[p] = next_standalone_id
            group_class[next_standalone_id] = c
            next_standalone_id += 1
            standalone += 1
            continue
        best_idx, best_dist = None, 10**6
        for (oi, _op, oh) in cands:
            d = h - oh
            if d < best_dist:
                best_dist = d; best_idx = oi
            if best_dist == 0:
                break
        if best_dist <= args.hamming_threshold:
            groups[p] = best_idx
            matched += 1
        else:
            groups[p] = next_standalone_id
            group_class[next_standalone_id] = c
            next_standalone_id += 1
            standalone += 1

    print(f"  matched to an original: {matched}")
    print(f"  standalone augmented (no near original): {standalone}")
    print(f"  total groups: {next_standalone_id}")

    # Build pool (Original + Augmented unioned), excluding nothing
    all_items = []  # (path, class_label_int, group_id)
    class_to_int = {c: i for i, c in enumerate(sorted(set(group_class.values())))}
    for p, c, _h in orig_hashes:
        all_items.append((p, class_to_int[c], groups[p]))
    for p, c, _h in aug_hashes:
        all_items.append((p, class_to_int[c], groups[p]))

    paths = np.array([x[0] for x in all_items])
    labels = np.array([x[1] for x in all_items])
    grps = np.array([x[2] for x in all_items])

    # Stratified-by-class, grouped-by-source split:
    # 1) Holdout test set: 15% by group (stratified on group majority label)
    # 2) From the remaining, build StratifiedGroupKFold folds for CV
    # 3) Also produce a single train/val cut from the pool for the "final" retrain
    rng = np.random.default_rng(args.seed)

    # build group -> (class_label, [item_indices])
    group_indices = defaultdict(list)
    for idx, g in enumerate(grps):
        group_indices[g].append(idx)
    group_ids = np.array(sorted(group_indices.keys()))
    group_labels = np.array([labels[group_indices[g][0]] for g in group_ids])

    # Stratified split of groups into (pool, test)
    pool_groups, test_groups = train_test_split(
        group_ids, test_size=args.test_frac, stratify=group_labels, random_state=args.seed
    )
    # From pool, further split val
    pool_labels = np.array([labels[group_indices[g][0]] for g in pool_groups])
    train_groups, val_groups = train_test_split(
        pool_groups, test_size=args.val_frac / (1 - args.test_frac),
        stratify=pool_labels, random_state=args.seed
    )

    def items_for(grps_subset):
        idxs = []
        for g in grps_subset:
            idxs.extend(group_indices[g])
        return [(paths[i], int(labels[i])) for i in idxs]

    splits = {
        "train": items_for(train_groups),
        "val":   items_for(val_groups),
        "test":  items_for(test_groups),
    }

    # k-fold over (train+val) groups
    pool_groups_sorted = np.concatenate([train_groups, val_groups])
    pool_labels_sorted = np.array([labels[group_indices[g][0]] for g in pool_groups_sorted])
    # Need also indices into the *pool_paths* list for the folds
    pool_items = items_for(pool_groups_sorted)
    pool_paths = [it[0] for it in pool_items]
    pool_labels_flat = [it[1] for it in pool_items]
    # And the group ID for every pool item
    pool_groups_flat = []
    for g in pool_groups_sorted:
        for _ in group_indices[g]:
            pool_groups_flat.append(int(g))

    sgkf = StratifiedGroupKFold(n_splits=args.n_folds, shuffle=True, random_state=args.seed)
    folds = []
    for fold_i, (tr_idx, va_idx) in enumerate(
        sgkf.split(np.zeros(len(pool_paths)), pool_labels_flat, groups=pool_groups_flat)
    ):
        folds.append({"train_idx": tr_idx.tolist(), "val_idx": va_idx.tolist()})

    out = {
        "seed": args.seed,
        "hamming_threshold": args.hamming_threshold,
        "classes": [c for c, _ in sorted(class_to_int.items(), key=lambda x: x[1])],
        "n_groups_total": int(next_standalone_id),
        "n_train_items": len(splits["train"]),
        "n_val_items": len(splits["val"]),
        "n_test_items": len(splits["test"]),
        "splits": splits,
        "pool_paths": pool_paths,
        "pool_labels": pool_labels_flat,
        "pool_groups": pool_groups_flat,
        "folds": folds,
    }
    with open(args.output, "w") as f:
        json.dump(out, f)
    print(f"\nManifest -> {args.output}")
    print(f"  train: {len(splits['train'])} items")
    print(f"  val:   {len(splits['val'])} items")
    print(f"  test:  {len(splits['test'])} items")
    print(f"  pool size for k-fold: {len(pool_paths)} items across {len(pool_groups_sorted)} groups")


if __name__ == "__main__":
    main()