src/preprocess.py

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