data/data_loader.py

import os
import glob
import random
from typing import List, Tuple, Dict

import numpy as np
import cv2
import torch
from torch.utils.data import Dataset, DataLoader, ConcatDataset
from skimage.metrics import structural_similarity as ssim_fn


def list_normal_images(root: str, category: str, dataset_name: str = "mvtec") -> List[str]:
    if dataset_name == "mvtec":
        pattern = os.path.join(root, category, "train", "good", "*.png")
    else:  # visa
        pattern = os.path.join(root, category, "train", "good", "*")
    files = sorted(glob.glob(pattern))
    return files


def list_test_images(root: str, category: str, dataset_name: str = "mvtec") -> List[Tuple[str, str, str]]:

    results = []
    test_root = os.path.join(root, category, "test")
    if not os.path.isdir(test_root):
        return results
    for defect_type in sorted(os.listdir(test_root)):
        img_paths = sorted(glob.glob(os.path.join(test_root, defect_type, "*")))
        for p in img_paths:
            if defect_type == "good":
                mask_path = None
            else:
                fname = os.path.splitext(os.path.basename(p))[0]
                mask_path = os.path.join(root, category, "ground_truth", defect_type, f"{fname}_mask.png")
                if not os.path.exists(mask_path):
                    mask_path = None
            results.append((p, defect_type, mask_path))
    return results


def list_dtd_textures(dtd_root: str) -> List[str]:
    pattern = os.path.join(dtd_root, "**", "*.jpg")
    files = sorted(glob.glob(pattern, recursive=True))
    return files



def load_and_resize(path: str, size: Tuple[int, int]) -> np.ndarray:

    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"Cannot read image: {path}")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    h, w = size
    img = cv2.resize(img, (w, h), interpolation=cv2.INTER_LINEAR)
    img = img.astype(np.float32) / 255.0
    return img


def load_mask(path: str, size: Tuple[int, int]) -> np.ndarray:
    m = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    h, w = size
    m = cv2.resize(m, (w, h), interpolation=cv2.INTER_NEAREST)
    m = (m > 127).astype(np.float32)
    return m


def compute_jnd_map(img: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY).astype(np.float32)

    bg_luminance = cv2.blur(gray, (9, 9))
    lum_jnd = np.where(
        bg_luminance <= 127,
        17.0 * (1 - np.sqrt(np.maximum(bg_luminance, 1) / 127.0)) + 3.0,
        3.0 / 128.0 * (bg_luminance - 127.0) + 3.0,
    )

    gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
    gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
    grad_mag = np.sqrt(gx ** 2 + gy ** 2)
    texture_jnd = 0.3 * grad_mag  

    jnd = lum_jnd + texture_jnd - 0.3 * np.minimum(lum_jnd, texture_jnd) 
    jnd = jnd / (jnd.max() + 1e-8) 
    return jnd.astype(np.float32)


PANELS = ["left", "right", "top", "bottom", "all"]


def get_panel_bbox(h: int, w: int, panel: str,
                    ratio_min: float = 0.5, ratio_max: float = 0.8) -> Tuple[int, int, int, int]:

    if panel == "all":
        return 0, h, 0, w

    rh = random.uniform(ratio_min, ratio_max)
    rw = random.uniform(ratio_min, ratio_max)
    region_h = int(h * rh)
    region_w = int(w * rw)

    if panel == "left":
        x0, x1 = 0, region_w
        y0, y1 = random.randint(0, h - region_h), None
        y1 = y0 + region_h
    elif panel == "right":
        x1, x0 = w, w - region_w
        y0 = random.randint(0, h - region_h)
        y1 = y0 + region_h
    elif panel == "top":
        y0, y1 = 0, region_h
        x0 = random.randint(0, w - region_w)
        x1 = x0 + region_w
    elif panel == "bottom":
        y1, y0 = h, h - region_h
        x0 = random.randint(0, w - region_w)
        x1 = x0 + region_w
    else:
        raise ValueError(f"Unknown panel: {panel}")

    return y0, y1, x0, x1


def sample_texture_source(normal_img: np.ndarray, dtd_paths: List[str], size: Tuple[int, int]) -> np.ndarray:
    """Anomaly texture T: DTD or augmented normal (resize, crop, flip, color-jitter)."""
    h, w = size
    if len(dtd_paths) > 0 and random.random() < 0.5:
        path = random.choice(dtd_paths)
        tex = load_and_resize(path, size)
    else:
        tex = normal_img.copy()

    # Random resize then crop back to (h, w)
    scale = random.uniform(0.8, 1.4)
    nh, nw = max(int(h * scale), h), max(int(w * scale), w)
    tex = cv2.resize(tex, (nw, nh), interpolation=cv2.INTER_LINEAR)
    y0 = random.randint(0, max(nh - h, 0))
    x0 = random.randint(0, max(nw - w, 0))
    tex = tex[y0:y0 + h, x0:x0 + w]
    if tex.shape[0] != h or tex.shape[1] != w:
        tex = cv2.resize(tex, (w, h), interpolation=cv2.INTER_LINEAR)

    if random.random() < 0.5:
        tex = np.fliplr(tex).copy()
    if random.random() < 0.5:
        tex = np.flipud(tex).copy()
    # Color jitter
    jitter = random.uniform(0.8, 1.2)
    brightness = random.uniform(-0.1, 0.1)
    tex = np.clip(tex * jitter + brightness, 0, 1)
    return tex.astype(np.float32)


def _smooth_weight(weight: np.ndarray, kernel_size: int) -> np.ndarray:
    """α_i as a smooth kernel (paper): softens anomaly/normal boundary. Easy → smaller kernel."""
    k = int(kernel_size)
    if k < 1:
        return weight
    if k % 2 == 0:
        k += 1
    return cv2.GaussianBlur(weight, (k, k), 0)


def synthesize_anomaly(normal_img: np.ndarray,
                        jnd_map: np.ndarray,
                        dtd_paths: List[str],
                        alpha_kernels: List[int] = (3, 7, 15),
                        beta: List[float] = (0.10, 0.05, 0.0),
                        area_ratio_min: float = 0.5,
                        area_ratio_max: float = 0.8,
                        ssim_thresholds: Tuple[float, float] = (0.9, 0.7),
                        # backward-compat alias used by older notebook cells
                        alpha: List = None) -> Dict[str, np.ndarray]:
    """
    Panel-guided synthesis (Eq. 1–2).
    Wi = αi * (J + βi) where αi is a smoothing kernel (not a scalar).
    A = mean_i (Wi*T + (1-Wi)*I)  — average keeps pixels in [0,1] (paper writes Σ).
    Severity GT refined via SSIM(I, A) into easy/medium/hard.
    """
    if alpha is not None:
        # If caller still passes scalar alphas, treat odds-ish values as kernel sizes
        alpha_kernels = [max(3, int(round(a * 15)) | 1) if isinstance(a, float) else int(a)
                         for a in alpha]

    h, w, _ = normal_img.shape
    panel = random.choice(PANELS)
    y0, y1, x0, x1 = get_panel_bbox(h, w, panel, area_ratio_min, area_ratio_max)

    texture = sample_texture_source(normal_img, dtd_paths, (h, w))

    J = jnd_map.astype(np.float32)
    A = normal_img.copy()
    mask = np.zeros((h, w), dtype=np.float32)

    region_slice = (slice(y0, y1), slice(x0, x1))
    I_region = normal_img[region_slice]
    T_region = texture[region_slice]
    J_region = J[region_slice]

    A_region_accum = np.zeros_like(I_region)
    for i in range(3):
        # Different random seeds / slight spatial jitter per level so regions may overlap
        noise = np.random.rand(*J_region.shape).astype(np.float32) * 0.02
        raw = np.clip(J_region + beta[i] + noise, 0.0, 1.0)
        Wi = _smooth_weight(raw, alpha_kernels[i])
        Wi = np.clip(Wi, 0.0, 1.0)[..., None]
        A_region_accum += Wi * T_region + (1.0 - Wi) * I_region
    A_region = A_region_accum / 3.0

    A[region_slice] = A_region
    mask[y0:y1, x0:x1] = 1.0

    # Refine & quantify into 3 defect levels via SSIM (paper Sec. 3.1)
    severity_mask = np.zeros((h, w), dtype=np.uint8)
    gray_orig = cv2.cvtColor((normal_img * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
    gray_syn = cv2.cvtColor((A * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
    _, ssim_map = ssim_fn(gray_orig, gray_syn, full=True)
    ssim_region = ssim_map[y0:y1, x0:x1]

    thr_easy, thr_hard = ssim_thresholds  # e.g. 0.9, 0.7
    sev_region = np.zeros_like(ssim_region, dtype=np.uint8)
    sev_region[ssim_region < thr_hard] = 3                          # hard
    sev_region[(ssim_region >= thr_hard) & (ssim_region < thr_easy)] = 2  # medium
    sev_region[(ssim_region >= thr_easy) & (ssim_region < 0.995)] = 1     # easy
    severity_mask[y0:y1, x0:x1] = sev_region

    return {
        "anomaly_image": A.astype(np.float32),
        "mask": mask,
        "severity_mask": severity_mask,
        "panel": panel,
    }



class OmniALDataset(Dataset):
    def __init__(self, root: str, categories: List[str], dtd_paths: List[str],
                 size: Tuple[int, int], dataset_name: str = "mvtec",
                 synthesis_cfg: dict = None):
        self.size = size
        self.dtd_paths = dtd_paths
        self.dataset_name = dataset_name
        syn = synthesis_cfg or {}
        self.alpha_kernels = syn.get("alpha_kernels", [3, 7, 15])
        self.beta = syn.get("beta", [0.10, 0.05, 0.0])
        self.area_ratio_min = syn.get("area_ratio_min", 0.5)
        self.area_ratio_max = syn.get("area_ratio_max", 0.8)
        thr = syn.get("ssim_thresholds", [0.9, 0.7])
        self.ssim_thresholds = (float(thr[0]), float(thr[1]))
        self.samples = []  # list of (path, category)
        for cat in categories:
            paths = list_normal_images(root, cat, dataset_name)
            self.samples += [(p, cat) for p in paths]
        if len(self.samples) == 0:
            print(f"[WARN] {dataset_name}: no images found under {root}")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        path, cat = self.samples[idx]
        img = load_and_resize(path, self.size)
        jnd = compute_jnd_map(img)

        syn = synthesize_anomaly(
            img, jnd, self.dtd_paths,
            alpha_kernels=self.alpha_kernels,
            beta=self.beta,
            area_ratio_min=self.area_ratio_min,
            area_ratio_max=self.area_ratio_max,
            ssim_thresholds=self.ssim_thresholds,
        )
        anomaly_img = syn["anomaly_image"]
        severity = syn["severity_mask"]

        to_tensor = lambda x: torch.from_numpy(x.transpose(2, 0, 1)).float() if x.ndim == 3 \
            else torch.from_numpy(x).float().unsqueeze(0)

        return {
            "image": to_tensor(img),
            "jnd": to_tensor(jnd),
            "anomaly_img": to_tensor(anomaly_img),
            "severity": torch.from_numpy(severity).long(),
            "category": cat,
        }



def build_dataloader(
    cfg: dict,
    dataset_name: str = "mvtec",
    shuffle: bool = True,
    categories: List[str] = None,
) -> DataLoader:
    """Build dataloader. Optional `categories` overrides cfg list (used for separate training)."""
    dtd_paths = list_dtd_textures(cfg["data"]["dtd_root"])

    if dataset_name == "mvtec":
        root = cfg["data"]["mvtec_root"]
        cats = categories if categories is not None else cfg["data"]["mvtec_categories"]
        size = tuple(cfg["data"]["mvtec_size"])
        batch_size = cfg["train"]["batch_size_mvtec"]
    else:
        root = cfg["data"]["visa_root"]
        cats = categories if categories is not None else cfg["data"]["visa_categories"]
        size = tuple(cfg["data"]["visa_size"])
        batch_size = cfg["train"]["batch_size_visa"]

    dataset = OmniALDataset(
        root, cats, dtd_paths, size, dataset_name,
        synthesis_cfg=cfg.get("synthesis", {}),
    )
    # Avoid empty loader when class has fewer samples than batch_size
    effective_bs = min(batch_size, max(len(dataset), 1))
    drop_last = len(dataset) >= effective_bs and effective_bs > 1
    loader = DataLoader(
        dataset,
        batch_size=effective_bs,
        shuffle=shuffle,
        num_workers=cfg["train"]["num_workers"],
        drop_last=drop_last,
        pin_memory=True,
    )
    return loader, dataset