Spaces:
Sleeping
Sleeping
| import os | |
| import glob | |
| import random | |
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| class GenImageDataset(Dataset): | |
| def __init__(self, root_dir, transform=None, sample_ratio=0.01, seed=42): | |
| """ | |
| Args: | |
| root_dir (str): Path to the root directory (e.g., 'genimage_test/test'). | |
| transform (callable, optional): Optional transform to be applied on a sample. | |
| sample_ratio (float): Ratio of data to sample (0.0 to 1.0). | |
| seed (int): Random seed for reproducibility. | |
| """ | |
| self.root_dir = root_dir | |
| self.transform = transform | |
| self.classes = ['ai'] | |
| self.samples = [] | |
| # Walk through the directory structure | |
| # Expected: root_dir/<generator>/[nature|ai]/<image> | |
| # Sometimes the structure is root_dir/test/<generator>/[nature|ai]/<image> | |
| # Get all generator folders | |
| if not os.path.exists(root_dir): | |
| raise FileNotFoundError(f"Directory not found: {root_dir}") | |
| # Auto-detect if there's a 'test' subdirectory (common in genimage datasets) | |
| test_dir = os.path.join(root_dir, 'test') | |
| if os.path.exists(test_dir) and os.path.isdir(test_dir): | |
| # Check if test_dir has generator folders | |
| test_contents = [d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))] | |
| if test_contents: | |
| print(f"Detected 'test' subdirectory, using {test_dir}") | |
| root_dir = test_dir | |
| generators = [d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))] | |
| print(f"Found {len(generators)} generator folders: {generators}") | |
| all_samples = [] | |
| for generator in generators: | |
| gen_path = os.path.join(root_dir, generator) | |
| for cls in self.classes: | |
| cls_path = os.path.join(gen_path, cls) | |
| if os.path.exists(cls_path): | |
| images = [] | |
| # Recursive search for image files | |
| extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff'] | |
| for ext in extensions: | |
| images.extend(glob.glob(os.path.join(cls_path, '**', ext), recursive=True)) | |
| # Also check for uppercase extensions | |
| images.extend(glob.glob(os.path.join(cls_path, '**', ext.upper()), recursive=True)) | |
| # Deduplicate to handle case-insensitive filesystems (Windows) | |
| images = sorted(list(set(images))) | |
| # Label: 0 for nature (real), 1 for ai (fake) | |
| label = 1 | |
| for img_path in images: | |
| all_samples.append((img_path, label)) | |
| # Sampling | |
| random.seed(seed) | |
| sample_size = int(len(all_samples) * sample_ratio) | |
| if sample_size > 0: | |
| self.samples = random.sample(all_samples, sample_size) | |
| else: | |
| self.samples = all_samples # Fallback if ratio is too small but we want something | |
| print(f"Loaded {len(self.samples)} samples from {len(all_samples)} total images.") | |
| def __len__(self): | |
| return len(self.samples) | |
| def __getitem__(self, idx): | |
| img_path, label = self.samples[idx] | |
| try: | |
| image = Image.open(img_path).convert('RGB') | |
| if self.transform: | |
| image = self.transform(image) | |
| return image, label, img_path | |
| except Exception as e: | |
| print(f"Error loading {img_path}: {e}") | |
| # Return a dummy or handle gracefully? For now, let's just skip or error. | |
| # In a real training loop, we might want to return None and use a collate_fn to filter. | |
| # For simplicity here, we'll raise. | |
| raise e | |