code/dataset.py

import os
from pathlib import Path
from typing import List

from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as TF


class PairedDataset(Dataset):
    """
    Paired image-to-image dataset for BBDM.

    Supported formats:
      - "side_by_side": single images with left=source, right=target
      - "trainA_trainB": root/{split}A/ and root/{split}B/ folders
      - "separate_domains": root/{domainA}/{split}/ and root/{domainB}/{split}/
        e.g. BCI dataset: root/HE/train/ (source) + root/IHC/train/ (target)
      - "auto": auto-detect from directory structure
    """

    def __init__(
        self,
        root: str = "data",
        split: str = "train",
        image_size: int = 64,
        augment: bool = False,
        format: str = "auto",
        source_domain: str = None,
        target_domain: str = None,
    ):
        """
        source_domain / target_domain: subdirectory names for "separate_domains" format.
            e.g. source_domain="HE", target_domain="IHC" for BCI dataset.
        """
        self.root = Path(root)
        self.split = split
        self.image_size = image_size
        self.augment = augment
        self.format = format

        # Resolve format and paths
        if format == "auto":
            if (self.root / split).exists():
                self.format = "side_by_side"
            elif (self.root / f"{split}A").exists() and (self.root / f"{split}B").exists():
                self.format = "trainA_trainB"
            else:
                # Auto-detect separate_domains: look for subdirs that contain split/
                subdirs = [d.name for d in self.root.iterdir() if d.is_dir() and (d / split).exists()]
                if len(subdirs) >= 2:
                    self.format = "separate_domains"
                    if source_domain is None or target_domain is None:
                        subdirs = sorted(subdirs)
                        source_domain = source_domain or subdirs[0]
                        target_domain = target_domain or subdirs[1]
                else:
                    self.format = "side_by_side"

        if self.format == "side_by_side":
            split_dir = self.root / split
            if not split_dir.exists():
                raise FileNotFoundError(f"Split directory not found: {split_dir}")
            self.paths = sorted(
                [p for p in split_dir.iterdir()
                 if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}]
            )
            self.pairs = None
        elif self.format == "separate_domains":
            if not source_domain or not target_domain:
                raise ValueError("source_domain and target_domain required for separate_domains format")
            dir_src = self.root / source_domain / split
            dir_tgt = self.root / target_domain / split
            if not dir_src.exists() or not dir_tgt.exists():
                raise FileNotFoundError(
                    f"Need {source_domain}/{split}/ and {target_domain}/{split}/ in {self.root}"
                )
            self.pairs = []
            for p in sorted(dir_src.iterdir()):
                if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}:
                    q = dir_tgt / p.name
                    if q.exists():
                        self.pairs.append((p, q))
            self.paths = None
        else:
            dir_a = self.root / f"{split}A"
            dir_b = self.root / f"{split}B"
            if not dir_a.exists() or not dir_b.exists():
                raise FileNotFoundError(f"Need {split}A/ and {split}B/ in {self.root}")
            self.pairs = []
            for p in dir_a.iterdir():
                if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}:
                    q = dir_b / p.name
                    if q.exists():
                        self.pairs.append((p, q))
            self.paths = None

        # Slightly upscale before crop for train-time spatial jitter.
        self.resize_for_crop = int(round(image_size * 1.125))
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # -> [-1, 1]

    def __len__(self) -> int:
        return len(self.paths) if self.paths is not None else len(self.pairs)

    def _load_pair(self, left, right):
        """Load and optionally augment a pair. left/right can be Path or PIL Image."""
        if isinstance(left, Path):
            with Image.open(left) as img:
                left = img.convert("RGB")
        if isinstance(right, Path):
            with Image.open(right) as img:
                right = img.convert("RGB")
        if self.augment:
            left, right = self._paired_train_transform(left, right)
        else:
            left = TF.resize(left, [self.image_size, self.image_size], interpolation=InterpolationMode.BICUBIC, antialias=True)
            right = TF.resize(right, [self.image_size, self.image_size], interpolation=InterpolationMode.BICUBIC, antialias=True)
        source = self.normalize(self.to_tensor(left))
        target = self.normalize(self.to_tensor(right))
        return source, target

    def __getitem__(self, idx: int):
        if self.format in ("trainA_trainB", "separate_domains"):
            left_path, right_path = self.pairs[idx]
            return self._load_pair(left_path, right_path)

        img_path = self.paths[idx]
        with Image.open(img_path) as img:
            img = img.convert("RGB")
            w, h = img.size
            if w % 2 != 0:
                raise ValueError(f"Expected even width for paired image, got width={w} in {img_path}")

            half_w = w // 2
            left = img.crop((0, 0, half_w, h))      # source
            right = img.crop((half_w, 0, w, h))     # target

        return self._load_pair(left, right)

    def _paired_train_transform(self, left: Image.Image, right: Image.Image):
        left = TF.resize(left, [self.resize_for_crop, self.resize_for_crop], interpolation=InterpolationMode.BICUBIC, antialias=True)
        right = TF.resize(right, [self.resize_for_crop, self.resize_for_crop], interpolation=InterpolationMode.BICUBIC, antialias=True)

        i, j, h, w = transforms.RandomCrop.get_params(left, output_size=(self.image_size, self.image_size))
        left = TF.crop(left, i, j, h, w)
        right = TF.crop(right, i, j, h, w)

        if torch.rand(1).item() < 0.5:
            left = TF.hflip(left)
            right = TF.hflip(right)

        return left, right


def denormalize(x: torch.Tensor) -> torch.Tensor:
    """Convert tensor range from [-1, 1] to [0, 1]."""
    return (x.clamp(-1, 1) + 1.0) * 0.5


# -----------------------------------------------------------------------------
# HuggingFace Datasets wrapper
# -----------------------------------------------------------------------------
# HFPairedDataset reads a paired image dataset from the HF Hub. The HF dataset
# is expected to have two image columns named via `source_col` / `target_col`
# (defaults: "source", "target"). It applies the same paired augmentation as
# PairedDataset so train.py can swap implementations without code changes.

class HFPairedDataset(Dataset):
    """Paired dataset backed by a HuggingFace Datasets repo.

    Args:
        repo_id: e.g. "augustander/bci-he2ihc"
        split: "train" / "test" / "val"
        image_size: target square crop size
        augment: 1.125x upscale → random crop → hflip (matches PairedDataset)
        source_col / target_col: column names in the HF dataset
        token: HF token string; if None, falls back to env HF_TOKEN
        cache_dir: optional override for HF cache location
    """

    def __init__(
        self,
        repo_id: str,
        split: str = "train",
        image_size: int = 256,
        augment: bool = False,
        source_col: str = "source",
        target_col: str = "target",
        token: str = None,
        cache_dir: str = None,
    ):
        from datasets import load_dataset  # local import — only required for HF mode

        self.image_size = image_size
        self.augment = augment
        self.source_col = source_col
        self.target_col = target_col

        self.ds = load_dataset(
            repo_id,
            split=split,
            token=token or os.environ.get("HF_TOKEN"),
            cache_dir=cache_dir or os.environ.get("HF_DATASETS_CACHE"),
        )

        self.resize_for_crop = int(round(image_size * 1.125))
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

    def __len__(self) -> int:
        return len(self.ds)

    def _aug(self, left: Image.Image, right: Image.Image):
        left = TF.resize(left, [self.resize_for_crop] * 2, interpolation=InterpolationMode.BICUBIC, antialias=True)
        right = TF.resize(right, [self.resize_for_crop] * 2, interpolation=InterpolationMode.BICUBIC, antialias=True)
        i, j, h, w = transforms.RandomCrop.get_params(left, output_size=(self.image_size, self.image_size))
        left = TF.crop(left, i, j, h, w)
        right = TF.crop(right, i, j, h, w)
        if torch.rand(1).item() < 0.5:
            left = TF.hflip(left)
            right = TF.hflip(right)
        return left, right

    def __getitem__(self, idx: int):
        row = self.ds[idx]
        left = row[self.source_col].convert("RGB")
        right = row[self.target_col].convert("RGB")
        if self.augment:
            left, right = self._aug(left, right)
        else:
            left = TF.resize(left, [self.image_size] * 2, interpolation=InterpolationMode.BICUBIC, antialias=True)
            right = TF.resize(right, [self.image_size] * 2, interpolation=InterpolationMode.BICUBIC, antialias=True)
        return self.normalize(self.to_tensor(left)), self.normalize(self.to_tensor(right))


def hf_download_checkpoint(spec: str, cache_dir: str = None) -> str:
    """Download a .pt checkpoint from the HF Hub.

    Args:
        spec: 'repo_id:filename' (e.g., 'augustander/bbdm-exp16:best_model.pt').
              If no ':filename' is given, tries common defaults in order.
        cache_dir: optional cache dir; falls back to HF default.

    Returns:
        Local path to the downloaded file.

    Reads HF_TOKEN from env for private repos.
    """
    from huggingface_hub import hf_hub_download

    if ":" in spec:
        repo_id, filename = spec.split(":", 1)
        candidates = [filename]
    else:
        repo_id = spec
        candidates = ["best_model.pt", "best.pt", "latest.pt"]

    last_err = None
    for fn in candidates:
        try:
            return hf_hub_download(
                repo_id=repo_id,
                filename=fn,
                token=os.environ.get("HF_TOKEN"),
                cache_dir=cache_dir or os.environ.get("HF_HOME"),
            )
        except Exception as e:
            last_err = e
    raise FileNotFoundError(
        f"None of {candidates} found in {repo_id} (HF). Last error: {last_err}"
    )


def build_dataset(
    *,
    hf_dataset: str = None,
    data_root: str = None,
    split: str = "train",
    image_size: int = 256,
    augment: bool = False,
    source_domain: str = None,
    target_domain: str = None,
    hf_source_col: str = "source",
    hf_target_col: str = "target",
):
    """Factory: returns HFPairedDataset if hf_dataset is set, else PairedDataset.

    Lets train.py swap between local and HF-backed data via a single arg.
    """
    if hf_dataset:
        return HFPairedDataset(
            repo_id=hf_dataset,
            split=split,
            image_size=image_size,
            augment=augment,
            source_col=hf_source_col,
            target_col=hf_target_col,
        )
    if data_root is None:
        raise ValueError("Either hf_dataset or data_root must be provided.")
    return PairedDataset(
        root=data_root,
        split=split,
        image_size=image_size,
        augment=augment,
        source_domain=source_domain,
        target_domain=target_domain,
    )