Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| from pathlib import Path | |
| from core.config import GENUINE_DIR, MASKS_DIR, MODEL_INPUT_SIZE, TAMPERED_DIR | |
| class TamperDataset(Dataset): | |
| def __init__(self, split: str = 'train', val_fraction: float = 0.15): | |
| genuine = sorted(GENUINE_DIR.glob('*.jpg')) + sorted(GENUINE_DIR.glob('*.png')) | |
| tampered = sorted(TAMPERED_DIR.glob('*.jpg')) + sorted(TAMPERED_DIR.glob('*.png')) | |
| all_items = [(p, 0) for p in genuine] + [(p, 1) for p in tampered] | |
| # Shuffle with fixed seed so train/val split is reproducible and class-balanced | |
| rng = np.random.default_rng(42) | |
| order = rng.permutation(len(all_items)).tolist() | |
| all_items = [all_items[i] for i in order] | |
| cut = int(len(all_items) * (1 - val_fraction)) | |
| self.items = all_items[:cut] if split == 'train' else all_items[cut:] | |
| self.size = MODEL_INPUT_SIZE | |
| self.augment = (split == 'train') | |
| def __len__(self): | |
| return len(self.items) | |
| def __getitem__(self, i) -> tuple[torch.Tensor, torch.Tensor, int]: | |
| path, label = self.items[i] | |
| img = cv2.imread(str(path), cv2.IMREAD_COLOR) | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| img = cv2.resize(img, (self.size, self.size)) | |
| img = img.astype(np.float32) / 255.0 | |
| if label == 1: | |
| mask_path = MASKS_DIR / (path.stem + '.png') | |
| if mask_path.exists(): | |
| mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) | |
| mask = cv2.resize(mask, (self.size, self.size)) | |
| mask = (mask > 127).astype(np.float32) | |
| else: | |
| mask = np.ones((self.size, self.size), dtype=np.float32) | |
| else: | |
| mask = np.zeros((self.size, self.size), dtype=np.float32) | |
| if self.augment: | |
| img, mask = _augment(img, mask) | |
| img_t = torch.from_numpy(img).permute(2, 0, 1) # (3, H, W) | |
| mask_t = torch.from_numpy(mask).unsqueeze(0) # (1, H, W) | |
| return img_t, mask_t, label | |
| def _augment(img: np.ndarray, mask: np.ndarray): | |
| # Flips | |
| if np.random.rand() > 0.5: | |
| img = img[:, ::-1, :].copy() | |
| mask = mask[:, ::-1].copy() | |
| if np.random.rand() > 0.5: | |
| img = img[::-1, :, :].copy() | |
| mask = mask[::-1, :].copy() | |
| # Brightness / contrast jitter | |
| if np.random.rand() > 0.5: | |
| alpha = np.random.uniform(0.7, 1.3) # contrast | |
| beta = np.random.uniform(-0.1, 0.1) # brightness | |
| img = np.clip(img * alpha + beta, 0, 1).astype(np.float32) | |
| # Gaussian noise | |
| if np.random.rand() > 0.5: | |
| noise = np.random.normal(0, 0.02, img.shape).astype(np.float32) | |
| img = np.clip(img + noise, 0, 1) | |
| return img, mask |