PcaCalculation.py

import os
import random
import numpy as np
from PIL import Image
import torch


def _collect_image_paths_recursive(root_dir: str, max_images: int | None = None):
    """
    Recursively collect image file paths under root_dir.
    Keeps the same behavior of limiting to max_images, but now across all subfolders.
    """
    exts = (".jpg", ".jpeg", ".png")
    paths = []

    for dirpath, _, filenames in os.walk(root_dir):
        for fname in filenames:
            if fname.lower().endswith(exts):
                paths.append(os.path.join(dirpath, fname))

    # Shuffle so we sample across folders/classes
    random.shuffle(paths)

    if max_images is not None:
        paths = paths[:max_images]

    return paths


def compute_pca_stream(
    train_dir: str,
    max_images: int = 1000,
    pixels_per_image: int = 100000,
    chunk_size: int = 20000
):
    """
    Streams pixel samples in chunks to compute mean and covariance without
    storing all samples in memory.

    Now scans ALL subfolders under train_dir (ImageFolder-style datasets).

    Args:
        train_dir: root directory of training images (may contain subfolders)
        max_images: max number of images to sample (after recursive collection + shuffle)
        pixels_per_image: total pixels to sample per image
        chunk_size: process pixels in chunks to bound memory use

    Returns:
        mean: (3,) array of per-channel means (in 0..255 scale)
        std: (3,) array of per-channel stddevs (in 0..255 scale)
        eigvals: (3,) PCA eigenvalues
        eigvecs: (3,3) PCA eigenvectors (columns correspond to eigvals)
    """
    # Accumulators
    sum_rgb = np.zeros(3, dtype=np.float64)       # Sum of RGB values
    sum_outer = np.zeros((3, 3), dtype=np.float64)  # Sum of outer products
    total = 0  # Total sampled pixels

    # Collect image paths recursively (do NOT change anything else about PCA math)
    image_paths = _collect_image_paths_recursive(train_dir, max_images=max_images)

    if len(image_paths) == 0:
        raise RuntimeError(f"No images found under: {train_dir}")

    # Process each image
    for path in image_paths:
        try:
            img = Image.open(path).convert("RGB")
        except Exception as e:
            print(f"[WARN] Skipping unreadable image: {path} ({e})")
            continue

        # IMPORTANT: keep full 0..255 scale (float)
        arr_img = np.asarray(img, np.float32)  # shape (H, W, 3), values 0..255
        if arr_img.ndim != 3 or arr_img.shape[2] != 3:
            print(f"[WARN] Skipping non-RGB image array: {path} shape={arr_img.shape}")
            continue

        H, W, _ = arr_img.shape

        remaining = pixels_per_image
        while remaining > 0:
            bs = min(chunk_size, remaining)

            ys = np.random.randint(0, H, size=bs)
            xs = np.random.randint(0, W, size=bs)

            samples = arr_img[ys, xs, :]  # (bs, 3)

            # Accumulate
            sum_rgb += samples.sum(axis=0)
            sum_outer += samples.T @ samples
            total += samples.shape[0]

            remaining -= bs

    if total < 2:
        raise RuntimeError("Not enough sampled pixels to compute covariance/PCA.")

    mean = sum_rgb / total
    cov = (sum_outer - total * np.outer(mean, mean)) / (total - 1)
    std = np.sqrt(np.diag(cov))

    eigvals, eigvecs = np.linalg.eigh(cov)
    idx = eigvals.argsort()[::-1]
    eigvals = eigvals[idx]
    eigvecs = eigvecs[:, idx]

    return (
        mean.astype(np.float32),
        std.astype(np.float32),
        eigvals.astype(np.float32),
        eigvecs.astype(np.float32),
    )


if __name__ == "__main__":
    # Root folder that contains all training images (including class subfolders)
    train_dir = "Test_256"

    print("Computing PCA...")

    mean, std, eigvals, eigvecs = compute_pca_stream(
        train_dir=train_dir,
        max_images=50000,
        pixels_per_image=100000,
        chunk_size=100000
    )

    print("Mean (R,G,B):", mean)
    print("Standard deviations (R,G,B):", std)
    print("Eigenvalues:", eigvals)
    print("First principal direction:", eigvecs[:, 0])

    pca_data = {
        "mean": mean,
        "std": std,
        "eigvals": eigvals,
        "eigvecs": eigvecs
    }

    torch.save(pca_data, "pca_values.pth")
    print("PCA values saved to pca_values.pth")