data.py

"""
Data loading for the Animal Faces (AFHQ) 32x32 image-classification dataset.

The Animal Faces dataset (AFHQ) contains ~16,130 32x32 RGB images in three
classes: **cat, dog, wild**. This module builds train/validation ``DataLoader``s
with standard augmentation and normalization, and can download the dataset if
it is not already present.

Expected on-disk layout (torchvision ``ImageFolder`` format):

    <data_dir>/afhq/train/<class>/*.png
    <data_dir>/afhq/val/<class>/*.png

Transforms
----------
- train : RandomRotation(10) -> RandomHorizontalFlip -> ToTensor -> Normalize
- val   : ToTensor -> Normalize
Normalization uses mean = std = (0.5, 0.5, 0.5), mapping pixels to [-1, 1].
"""

import os
import pathlib

import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder

# Public mirror used in the course materials.
AFHQ_URL = "https://osf.io/kgfvj/download"
MEAN = (0.5, 0.5, 0.5)
STD = (0.5, 0.5, 0.5)


def download_afhq(data_dir: str) -> pathlib.Path:
    """Download and unzip the AFHQ 32x32 dataset into ``data_dir`` if missing.

    Returns the path to the ``afhq`` folder (containing ``train`` and ``val``).
    """
    import requests
    from zipfile import ZipFile

    root = pathlib.Path(data_dir)
    root.mkdir(parents=True, exist_ok=True)
    afhq_dir = root / "afhq"
    if afhq_dir.exists():
        return afhq_dir

    zip_path = root / "AnimalFaces32x32.zip"
    print("Downloading AnimalFaces (AFHQ) dataset...")
    r = requests.get(AFHQ_URL, allow_redirects=True, timeout=120)
    zip_path.write_bytes(r.content)
    with ZipFile(zip_path, "r") as zf:
        zf.extractall(root)
    zip_path.unlink(missing_ok=True)
    # The archive extracts to a folder that contains 'afhq'; normalize the path.
    if not afhq_dir.exists():
        for p in root.rglob("afhq"):
            return p
    return afhq_dir


def _seed_worker(worker_id: int) -> None:
    worker_seed = torch.initial_seed() % 2 ** 32
    import numpy as np
    import random
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def get_dataloaders(data_dir: str = "./data", batch_size: int = 64,
                    seed: int = 42, download: bool = False, num_workers: int = 2):
    """Build (train_loader, val_loader, classes) for AFHQ.

    Parameters
    ----------
    data_dir : str
        Directory containing (or to receive) the ``afhq`` folder.
    batch_size : int
    seed : int
        Seed for the DataLoader shuffling generator.
    download : bool
        If True, download AFHQ when it is not already present.
    num_workers : int

    Returns
    -------
    train_loader, val_loader : torch.utils.data.DataLoader
    classes : list[str]
    """
    afhq_dir = pathlib.Path(data_dir) / "afhq"
    if download and not afhq_dir.exists():
        afhq_dir = download_afhq(data_dir)
    if not afhq_dir.exists():
        raise FileNotFoundError(
            f"AFHQ not found at {afhq_dir}. Pass download=True or place the "
            f"dataset there as afhq/train/<class> and afhq/val/<class>."
        )

    train_tf = transforms.Compose([
        transforms.RandomRotation(10),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),
    ])
    val_tf = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),
    ])

    train_ds = ImageFolder(afhq_dir / "train", transform=train_tf)
    val_ds = ImageFolder(afhq_dir / "val", transform=val_tf)

    g = torch.Generator()
    g.manual_seed(seed)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                              num_workers=num_workers, worker_init_fn=_seed_worker, generator=g)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
                            num_workers=num_workers, worker_init_fn=_seed_worker, generator=g)
    return train_loader, val_loader, train_ds.classes