""" 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()