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, )