import os import cv2 import numpy as np import torch import albumentations as A from torch.utils.data import Dataset, DataLoader from albumentations.pytorch import ToTensorV2 class MVTecDefectDataset(Dataset): def __init__(self, root_dir, object_class, split="train", transform=None, defect_type=None, dilate_mask=False, mask_kernel_size=3): """ Args: root_dir (str): Directory with MVTec AD dataset object_class (str): Object class (e.g., 'bottle', 'cable', etc.) split (str): 'train' or 'test' transform: Optional transform to be applied defect_type (str): Specific defect type to load (e.g., 'broken_large'). If None, loads all defect types. """ # Ensure path uses correct operating system format self.root_dir = os.path.normpath(root_dir) self.object_class = object_class self.split = split self.transform = transform self.target_defect_type = defect_type self.dilate_mask = dilate_mask self.mask_kernel_size = mask_kernel_size print(f"Initializing Dataset: root_dir={self.root_dir}, object_class={object_class}, split={split}") if defect_type: print(f"Target defect type: {defect_type}") # Identify defect types self.defect_types = [] if split == "train": defect_path = os.path.join(self.root_dir, object_class, "train", "defective") print(f"Searching for defect types in: {defect_path}") if os.path.exists(defect_path): all_defect_types = [d for d in os.listdir(defect_path) if os.path.isdir(os.path.join(defect_path, d))] # If a specific type is requested, only use that if defect_type and defect_type in all_defect_types: self.defect_types = [defect_type] print(f"Using specified defect type: {self.defect_types}") elif defect_type: print(f"Warning: Requested defect type '{defect_type}' not found. Available: {all_defect_types}") self.defect_types = all_defect_types else: self.defect_types = all_defect_types print(f"Found defect types: {self.defect_types}") else: print(f"Warning: Directory does not exist {defect_path}") # Debugging helper: Check parent directory parent_dir = os.path.dirname(defect_path) if os.path.exists(parent_dir): print(f"Parent directory {parent_dir} exists, containing: {os.listdir(parent_dir)}") if os.path.exists(self.root_dir): root_contents = os.listdir(self.root_dir) print(f"Root {self.root_dir} exists, containing: {root_contents[:5]}... ({len(root_contents)} items total)") # Load image and mask paths self.images = [] self.masks = [] # Handle Test Split (only loads the 'good' directory for anomaly detection baselines) if split == "test": good_dir = os.path.join(self.root_dir, object_class, "test", "good") print(f"Loading test set from: {good_dir}") if os.path.exists(good_dir): good_files = sorted([f for f in os.listdir(good_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]) print(f"Found {len(good_files)} 'good' samples in test set") for good_file in good_files: self.images.append(os.path.join(good_dir, good_file)) # Masks are generated randomly for good images during training/inference self.masks.append(None) else: print(f"Warning: Test directory not found {good_dir}") # Load Defect Images (Train Split) else: for defect_type in self.defect_types: img_dir = os.path.join(self.root_dir, object_class, "train", "defective", defect_type) mask_dir = os.path.join(self.root_dir, object_class, "train", "defective_masks", defect_type) print(f"Processing defect: {defect_type}") print(f" Images: {img_dir}") print(f" Masks: {mask_dir}") if os.path.exists(img_dir) and os.path.exists(mask_dir): img_files = sorted([f for f in os.listdir(img_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]) print(f" Found {len(img_files)} image files") matched = 0 for img_file in img_files: img_path = os.path.join(img_dir, img_file) # Convention: mask is base_name + _mask.png base_name = os.path.splitext(img_file)[0] mask_file = f"{base_name}_mask.png" mask_path = os.path.join(mask_dir, mask_file) # Fallback for alternative mask naming patterns if not os.path.exists(mask_path): possible_masks = [f for f in os.listdir(mask_dir) if base_name in f] if possible_masks: mask_path = os.path.join(mask_dir, possible_masks[0]) else: print(f" Warning: No mask found for {img_file}, skipping") continue self.images.append(img_path) self.masks.append(mask_path) matched += 1 print(f" Successfully paired {matched} image-mask sets") print(f"Total loaded: {len(self.images)} {split} images") if len(self.images) == 0: print("Warning: Dataset is empty!") def __len__(self): return len(self.images) def generate_random_mask(self, image_size): """Generate random rectangular masks for object loss (integrity learning)""" mask = np.zeros(image_size, dtype=np.float32) num_rectangles = 30 h, w = image_size for _ in range(num_rectangles): # Rectangle size between 3% and 25% of image dimensions min_size = int(min(h, w) * 0.03) max_size = int(min(h, w) * 0.25) rect_h = np.random.randint(min_size, max_size) rect_w = np.random.randint(min_size, max_size) y = np.random.randint(0, h - rect_h) x = np.random.randint(0, w - rect_w) mask[y:y+rect_h, x:x+rect_w] = 1.0 return mask def __getitem__(self, idx): img_path = self.images[idx] mask_path = self.masks[idx] # Load image (OpenCV loads BGR, convert to RGB) image = cv2.imread(img_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Load ground truth mask or generate a random one if mask_path is not None: mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # Dilate mask (only if enabled) if self.dilate_mask: # Ensure kernel size is odd, otherwise OpenCV crashes k_size = self.mask_kernel_size if self.mask_kernel_size % 2 == 1 else self.mask_kernel_size + 1 kernel = np.ones((k_size, k_size), np.uint8) mask = cv2.dilate(mask, kernel, iterations=1) mask = mask.astype(np.float32) / 255.0 # Normalize to [0, 1] else: mask = self.generate_random_mask((image.shape[0], image.shape[1])) # --- NEW CODE START: SMART RESIZING --- h, w = image.shape[:2] target_size = 512 # Case 1: Image is too small (e.g., 400x400) -> Upscale it if h < target_size or w < target_size: # Use INTER_CUBIC to keep lines as sharp as possible image = cv2.resize(image, (target_size, target_size), interpolation=cv2.INTER_CUBIC) if mask is not None: # Use NEAREST for masks to avoid creating gray pixels at edges mask = cv2.resize(mask, (target_size, target_size), interpolation=cv2.INTER_NEAREST) # Case 2: Image is large (1024x1024) -> Do nothing here! # The RandomCrop in the transform will handle it, preserving full detail. # --- NEW CODE END --- # Apply Albumentations transformations if self.transform: augmented = self.transform(image=image, mask=mask) image = augmented['image'] mask = augmented['mask'] # Create background (masked image) for the inpainting input: I * (1 - M) background = image * (1 - mask) # Adjusted mask used for Object Loss calculation adjusted_mask = mask + 0.3 * (1 - mask) if mask_path is None else mask return { 'image': image, 'mask': mask, 'background': background, 'adjusted_mask': adjusted_mask, 'is_defect': mask_path is not None, 'object_class': self.object_class } def get_data_loaders(root_dir, object_class, batch_size=4, defect_type=None, dilate_mask=False, mask_kernel_size=3): """Creates training and testing DataLoaders with preprocessing pipelines""" # Training pipeline: Includes random scaling for better generalization train_transform = A.Compose([ A.RandomScale(scale_limit=(0.0, 0.125), p=1.0), # Random scale between 1.0 and 1.125 A.RandomCrop(height=512, width=512), A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ToTensorV2() ], additional_targets={'mask': 'mask', 'background': 'image', 'adjusted_mask': 'mask'}) # Test pipeline: Simple resize and normalize test_transform = A.Compose([ A.Resize(512, 512), A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ToTensorV2() ], additional_targets={'mask': 'mask', 'background': 'image', 'adjusted_mask': 'mask'}) train_dataset = MVTecDefectDataset( root_dir=root_dir, object_class=object_class, split="train", transform=train_transform, defect_type=defect_type, dilate_mask=dilate_mask, mask_kernel_size=mask_kernel_size ) test_dataset = MVTecDefectDataset( root_dir=root_dir, object_class=object_class, split="test", transform=test_transform, dilate_mask=dilate_mask, mask_kernel_size=mask_kernel_size ) train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True ) test_loader = DataLoader( test_dataset, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True ) return train_loader, test_loader