Spaces:
Sleeping
Sleeping
| """ | |
| Dataset loader for AI Image Detection | |
| """ | |
| import os | |
| from pathlib import Path | |
| import torch | |
| from torch.utils.data import Dataset | |
| from torchvision import transforms | |
| from PIL import Image | |
| class AIImageDataset(Dataset): | |
| """Custom Dataset for AI-generated vs Natural images""" | |
| def __init__(self, root_dir, transform=None, class_names=['Natural', 'Synthetic']): | |
| """ | |
| Args: | |
| root_dir (str): Root directory containing image class folders | |
| transform (callable, optional): Optional transform to be applied on images | |
| class_names (list): List of class folder names | |
| """ | |
| self.root_dir = Path(root_dir) | |
| self.transform = transform | |
| self.class_names = class_names | |
| self.images = [] | |
| self.labels = [] | |
| # Load image paths and labels, validating images and skipping corrupted files | |
| skipped = 0 | |
| for class_idx, class_name in enumerate(class_names): | |
| class_dir = self.root_dir / class_name | |
| if not class_dir.exists(): | |
| print(f"Warning: {class_dir} does not exist") | |
| continue | |
| for img_path in class_dir.glob('*'): | |
| if img_path.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.gif']: | |
| try: | |
| # Validate image file quickly by opening and verifying header | |
| with Image.open(img_path) as im: | |
| im.verify() | |
| self.images.append(img_path) | |
| self.labels.append(class_idx) | |
| except Exception as e: | |
| skipped += 1 | |
| print(f"Skipped corrupt image {img_path}: {type(e).__name__}: {e}") | |
| print(f"Loaded {len(self.images)} images from {root_dir} (skipped {skipped} corrupt files)") | |
| for class_idx, class_name in enumerate(class_names): | |
| count = sum(1 for label in self.labels if label == class_idx) | |
| print(f" {class_name}: {count}") | |
| def __len__(self): | |
| return len(self.images) | |
| def __getitem__(self, idx): | |
| img_path = self.images[idx] | |
| label = self.labels[idx] | |
| try: | |
| image = Image.open(img_path).convert('RGB') | |
| except Exception as e: | |
| print(f"Error loading image {img_path}: {type(e).__name__}: {e}") | |
| # Return a dummy image in case of error to keep DataLoader stable | |
| image = Image.new('RGB', (224, 224)) | |
| if self.transform: | |
| image = self.transform(image) | |
| return image, label | |
| def get_transforms(image_size=384, mode='train', normalize_mean=None, normalize_std=None): | |
| """ | |
| Get image transforms for preprocessing | |
| Args: | |
| image_size (int): Target image size | |
| mode (str): 'train' for augmentation, 'val'/'test' for no augmentation | |
| normalize_mean (list): Normalization mean values | |
| normalize_std (list): Normalization std values | |
| Returns: | |
| transforms.Compose: Composed transforms | |
| """ | |
| if normalize_mean is None: | |
| normalize_mean = [0.485, 0.456, 0.406] | |
| if normalize_std is None: | |
| normalize_std = [0.229, 0.224, 0.225] | |
| if mode == 'train': | |
| return transforms.Compose([ | |
| transforms.RandomResizedCrop(image_size, scale=(0.8, 1.0)), | |
| transforms.RandomHorizontalFlip(p=0.5), | |
| transforms.RandomVerticalFlip(p=0.2), | |
| transforms.RandomRotation(15), | |
| transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), | |
| transforms.RandomAffine(degrees=10, translate=(0.1, 0.1)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=normalize_mean, std=normalize_std) | |
| ]) | |
| else: | |
| return transforms.Compose([ | |
| transforms.Resize((image_size, image_size)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=normalize_mean, std=normalize_std) | |
| ]) | |