omnial / data /data_loader.py

farsi comments

04ec2fb 2 months ago

11.9 kB

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

	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 (WiT + (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