Class TinyImageNetDataset

ML/src/python/neuralforge/data/datasets.py:166–208

        return self.dataset[idx]

class TinyImageNetDataset:
    def __init__(self, root='./data', train=True, transform=None, download=True):
        if transform is None:
            if train:
                transform = transforms.Compose([
                    transforms.RandomCrop(64, padding=8),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                ])
            else:
                transform = transforms.Compose([
                    transforms.ToTensor(),
                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                ])
        
        import zipfile
        import urllib.request
        
        data_dir = os.path.join(root, 'tiny-imagenet-200')
        if download and not os.path.exists(data_dir):
            print("Downloading Tiny ImageNet (237 MB)...")
            url = 'http://cs231n.stanford.edu/tiny-imagenet-200.zip'
            zip_path = os.path.join(root, 'tiny-imagenet-200.zip')
            
            try:
                urllib.request.urlretrieve(url, zip_path)
                print("Extracting...")
                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                    zip_ref.extractall(root)
                os.remove(zip_path)
            except Exception as e:
                print(f"Download failed: {e}")
                print("Please download manually from: http://cs231n.stanford.edu/tiny-imagenet-200.zip")
        
        split = 'train' if train else 'val'
        self.dataset = datasets.ImageFolder(os.path.join(data_dir, split), transform=transform)
    
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, idx):
        return self.dataset[idx]

class Food101Dataset:
    def __init__(self, root='./data', split='train', transform=None, download=True):