import random import shutil from pathlib import Path import numpy as np import matplotlib.pyplot as plt from sklearn.model_selection import train_test_split from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array from tensorflow.keras.applications.mobilenet_v2 import preprocess_input CLASSES = ["with_mask", "without_mask"] IMG_SIZE = (224, 224) BATCH_SIZE = 32 def split_dataset( src_dir: Path, dest_dir: Path, val_split: float = 0.15, test_split: float = 0.15, seed: int = 42, overwrite: bool = False, ) -> dict: """ Copies images from src_dir/{class}/ into dest_dir/train|val|test/{class}/. Splits are stratified per class. Returns a summary dict: {class: {split: count}}. """ src_dir = Path(src_dir) dest_dir = Path(dest_dir) if dest_dir.exists() and any(dest_dir.iterdir()) and not overwrite: print(f"[INFO] Split directory already exists at {dest_dir}. Pass overwrite=True to redo.") return _count_existing_split(dest_dir) summary = {} for cls in CLASSES: cls_dir = src_dir / cls images = sorted( list(cls_dir.glob("*.jpg")) + list(cls_dir.glob("*.jpeg")) + list(cls_dir.glob("*.png")) ) if not images: print(f"[WARNING] No images found in {cls_dir}") continue train_val, test = train_test_split(images, test_size=test_split, random_state=seed) val_ratio = val_split / (1.0 - test_split) train, val = train_test_split(train_val, test_size=val_ratio, random_state=seed) summary[cls] = {"train": len(train), "val": len(val), "test": len(test)} for split_name, split_files in [("train", train), ("val", val), ("test", test)]: out_dir = dest_dir / split_name / cls out_dir.mkdir(parents=True, exist_ok=True) for f in split_files: shutil.copy2(f, out_dir / f.name) print("[INFO] Dataset split complete.") _print_summary(summary) return summary def _count_existing_split(dest_dir: Path) -> dict: summary = {} for split in ("train", "val", "test"): for cls in CLASSES: d = dest_dir / split / cls count = len(list(d.glob("*.*"))) if d.exists() else 0 summary.setdefault(cls, {})[split] = count _print_summary(summary) return summary def _print_summary(summary: dict) -> None: print(f"\n{'Class':<22} {'Train':>7} {'Val':>7} {'Test':>7} {'Total':>8}") print("-" * 52) totals = {"train": 0, "val": 0, "test": 0} for cls, counts in summary.items(): t = counts.get("train", 0) v = counts.get("val", 0) s = counts.get("test", 0) print(f" {cls:<20} {t:>7} {v:>7} {s:>7} {t+v+s:>8}") totals["train"] += t; totals["val"] += v; totals["test"] += s print("-" * 52) total = sum(totals.values()) print(f" {'TOTAL':<20} {totals['train']:>7} {totals['val']:>7} {totals['test']:>7} {total:>8}") print(f"\n Split ratios → " f"train={totals['train']/total*100:.1f}% " f"val={totals['val']/total*100:.1f}% " f"test={totals['test']/total*100:.1f}%\n") # ── Augmentation configs ────────────────────────────────────────────────────── def get_train_augmenter() -> ImageDataGenerator: return ImageDataGenerator( preprocessing_function=preprocess_input, rotation_range=20, zoom_range=0.10, width_shift_range=0.10, height_shift_range=0.10, shear_range=0.10, horizontal_flip=True, brightness_range=[0.8, 1.2], fill_mode="nearest", ) def get_eval_augmenter() -> ImageDataGenerator: return ImageDataGenerator(preprocessing_function=preprocess_input) # ── Generator pipeline ──────────────────────────────────────────────────────── def build_generators( split_dir: Path, batch_size: int = BATCH_SIZE, img_size: tuple = IMG_SIZE, ): """ Returns (train_gen, val_gen, test_gen, class_indices). train_gen applies augmentation; val/test generators only preprocess. """ split_dir = Path(split_dir) train_gen = get_train_augmenter().flow_from_directory( str(split_dir / "train"), target_size=img_size, batch_size=batch_size, class_mode="categorical", shuffle=True, seed=42, ) val_gen = get_eval_augmenter().flow_from_directory( str(split_dir / "val"), target_size=img_size, batch_size=batch_size, class_mode="categorical", shuffle=False, ) test_gen = get_eval_augmenter().flow_from_directory( str(split_dir / "test"), target_size=img_size, batch_size=batch_size, class_mode="categorical", shuffle=False, ) return train_gen, val_gen, test_gen, train_gen.class_indices # ── Visualization helpers ───────────────────────────────────────────────────── def visualize_augmentations(img_path: str, n: int = 8, seed: int = 0) -> None: """Show n augmented versions of a single image side-by-side.""" aug = ImageDataGenerator( rotation_range=20, zoom_range=0.10, width_shift_range=0.10, height_shift_range=0.10, shear_range=0.10, horizontal_flip=True, brightness_range=[0.8, 1.2], fill_mode="nearest", ) img = load_img(img_path, target_size=IMG_SIZE) img_arr = img_to_array(img) img_arr = np.expand_dims(img_arr, axis=0) cols = 4 rows = (n + cols) // cols # +1 col for original fig, axes = plt.subplots(rows, cols, figsize=(14, rows * 3.5)) axes = axes.flatten() # First cell = original axes[0].imshow(img) axes[0].set_title("Original", fontweight="bold", color="green") axes[0].axis("off") gen = aug.flow(img_arr, batch_size=1, seed=seed) for i in range(1, n + 1): aug_img = next(gen)[0].astype("uint8") axes[i].imshow(aug_img) axes[i].set_title(f"Aug #{i}") axes[i].axis("off") for j in range(n + 1, len(axes)): axes[j].axis("off") plt.suptitle("Data Augmentation Samples", fontsize=14, fontweight="bold") plt.tight_layout() plt.show() def visualize_pixel_distribution(img_path: str) -> None: """Compare raw vs preprocess_input pixel value distributions.""" img = load_img(img_path, target_size=IMG_SIZE) raw = img_to_array(img) # [0, 255] norm = preprocess_input(raw.copy()) # [-1, 1] fig, axes = plt.subplots(1, 2, figsize=(13, 4)) colors = ["#e74c3c", "#2ecc71", "#3498db"] labels = ["Red", "Green", "Blue"] for ch, (col, lbl) in enumerate(zip(colors, labels)): axes[0].hist(raw[..., ch].flatten(), bins=50, color=col, alpha=0.55, label=lbl) axes[1].hist(norm[..., ch].flatten(), bins=50, color=col, alpha=0.55, label=lbl) axes[0].set_title("Raw pixel values [0 – 255]") axes[0].set_xlabel("Pixel value") axes[0].legend() axes[1].set_title("After preprocess_input [-1 – 1]") axes[1].set_xlabel("Pixel value") axes[1].legend() plt.suptitle("Pixel Value Distribution Before / After Preprocessing", fontsize=13, fontweight="bold") plt.tight_layout() plt.show() def show_batch(generator, class_indices: dict, n: int = 8) -> None: """Display one batch of images from a generator with their labels.""" class_names = {v: k for k, v in class_indices.items()} images, labels = next(generator) cols = 4 rows = (n + cols - 1) // cols fig, axes = plt.subplots(rows, cols, figsize=(14, rows * 3.5)) axes = axes.flatten() for i in range(min(n, len(images))): # Reverse preprocess_input for display: [-1,1] → [0,255] display = (images[i] + 1.0) / 2.0 display = np.clip(display, 0, 1) label = class_names[np.argmax(labels[i])] color = "green" if label == "with_mask" else "red" axes[i].imshow(display) axes[i].set_title(label, color=color, fontsize=10) axes[i].axis("off") for j in range(n, len(axes)): axes[j].axis("off") plt.suptitle("Sample Training Batch (after augmentation + preprocessing)", fontsize=13, fontweight="bold") plt.tight_layout() plt.show()