landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/fid.py

@@ -0,0 +1,244 @@
+"""Self-contained FID computation using InceptionV3 feature extraction.
+
+Avoids dependency on torch-fidelity by implementing FID directly.
+Supports GPU acceleration, batched processing, and caching.
+
+Usage:
+    from landmarkdiff.fid import compute_fid_from_dirs, compute_fid_from_arrays
+
+    # From directories
+    fid = compute_fid_from_dirs("path/to/real", "path/to/generated")
+
+    # From numpy arrays
+    fid = compute_fid_from_arrays(real_images, generated_images)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+    import torch
+    import torch.nn as nn
+    from torch.utils.data import DataLoader, Dataset
+
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+    Dataset = object  # type: ignore[misc,assignment]
+
+
+def _load_inception_v3() -> Any:
+    """Load InceptionV3 with pool3 features (2048-dim)."""
+    from torchvision.models import Inception_V3_Weights, inception_v3
+
+    model = inception_v3(weights=Inception_V3_Weights.IMAGENET1K_V1)
+    # We want features from the avg pool layer (2048-dim)
+    # Remove the final FC layer
+    model.fc = nn.Identity()
+    model.eval()
+    return model
+
+
+class ImageFolderDataset(Dataset):
+    """Simple dataset that loads images from a directory."""
+
+    def __init__(self, directory: str | Path, image_size: int = 299):
+        self.directory = Path(directory)
+        exts = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+        self.files = sorted(
+            f for f in self.directory.iterdir() if f.suffix.lower() in exts and f.is_file()
+        )
+        self.image_size = image_size
+
+    def __len__(self) -> int:
+        return len(self.files)
+
+    def __getitem__(self, idx: int) -> Any:
+        import cv2
+
+        img = cv2.imread(str(self.files[idx]))
+        if img is None:
+            # Return zeros if image can't be loaded
+            return torch.zeros(3, self.image_size, self.image_size)
+        img = cv2.resize(img, (self.image_size, self.image_size))
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        # Normalize to [0, 1] then ImageNet normalize
+        t = torch.from_numpy(img.astype(np.float32) / 255.0).permute(2, 0, 1)
+        t = _imagenet_normalize(t)
+        return t
+
+
+class NumpyArrayDataset(Dataset):
+    """Dataset wrapping a list of numpy arrays."""
+
+    def __init__(self, images: list[np.ndarray], image_size: int = 299):
+        self.images = images
+        self.image_size = image_size
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def __getitem__(self, idx: int) -> Any:
+        import cv2
+
+        img = self.images[idx]
+        if img.shape[:2] != (self.image_size, self.image_size):
+            img = cv2.resize(img, (self.image_size, self.image_size))
+        if img.ndim == 2:
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+        elif img.shape[2] == 4:
+            img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+        elif img.shape[2] == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        t = torch.from_numpy(img.astype(np.float32) / 255.0).permute(2, 0, 1)
+        t = _imagenet_normalize(t)
+        return t
+
+
+def _imagenet_normalize(t: torch.Tensor) -> torch.Tensor:
+    """Apply ImageNet normalization."""
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+    return (t - mean) / std
+
+
+def _extract_features(
+    model: nn.Module,
+    dataloader: DataLoader,
+    device: torch.device,
+) -> np.ndarray:
+    """Extract InceptionV3 pool3 features from a dataloader."""
+    features = []
+    with torch.no_grad():
+        for batch in dataloader:
+            batch = batch.to(device)
+            feat = model(batch)
+            if isinstance(feat, tuple):
+                feat = feat[0]
+            features.append(feat.cpu().numpy())
+    return np.concatenate(features, axis=0)
+
+
+def _compute_statistics(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Compute mean and covariance of feature vectors."""
+    if features.shape[0] < 2:
+        raise ValueError(f"FID requires at least 2 images, got {features.shape[0]}")
+    mu = np.mean(features, axis=0)
+    sigma = np.cov(features, rowvar=False)
+    return mu, sigma
+
+
+def _calculate_fid(
+    mu1: np.ndarray,
+    sigma1: np.ndarray,
+    mu2: np.ndarray,
+    sigma2: np.ndarray,
+) -> float:
+    """Calculate FID given two sets of statistics.
+
+    FID = ||mu1 - mu2||^2 + Tr(sigma1 + sigma2 - 2*sqrt(sigma1*sigma2))
+    """
+    from scipy.linalg import sqrtm
+
+    diff = mu1 - mu2
+    covmean = sqrtm(sigma1 @ sigma2)
+
+    # Handle numerical instability
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
+
+    fid = diff @ diff + np.trace(sigma1 + sigma2 - 2 * covmean)
+    return float(max(fid, 0.0))
+
+
+def compute_fid_from_dirs(
+    real_dir: str | Path,
+    generated_dir: str | Path,
+    batch_size: int = 32,
+    num_workers: int = 4,
+    device: str | None = None,
+) -> float:
+    """Compute FID between two directories of images.
+
+    Args:
+        real_dir: Path to real images.
+        generated_dir: Path to generated images.
+        batch_size: Batch size for feature extraction.
+        num_workers: DataLoader workers.
+        device: "cuda" or "cpu". Auto-detects if None.
+
+    Returns:
+        FID score (lower = better).
+    """
+    if not HAS_TORCH:
+        raise ImportError("PyTorch required for FID computation")
+
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    dev = torch.device(device)
+
+    model = _load_inception_v3().to(dev)
+
+    real_ds = ImageFolderDataset(real_dir)
+    gen_ds = ImageFolderDataset(generated_dir)
+
+    if len(real_ds) == 0 or len(gen_ds) == 0:
+        raise ValueError("Need at least 1 image in each directory")
+
+    real_loader = DataLoader(
+        real_ds, batch_size=batch_size, num_workers=num_workers, pin_memory=True
+    )
+    gen_loader = DataLoader(gen_ds, batch_size=batch_size, num_workers=num_workers, pin_memory=True)
+
+    real_features = _extract_features(model, real_loader, dev)
+    gen_features = _extract_features(model, gen_loader, dev)
+
+    mu_real, sigma_real = _compute_statistics(real_features)
+    mu_gen, sigma_gen = _compute_statistics(gen_features)
+
+    return _calculate_fid(mu_real, sigma_real, mu_gen, sigma_gen)
+
+
+def compute_fid_from_arrays(
+    real_images: list[np.ndarray],
+    generated_images: list[np.ndarray],
+    batch_size: int = 32,
+    device: str | None = None,
+) -> float:
+    """Compute FID from lists of numpy arrays.
+
+    Args:
+        real_images: List of (H, W, 3) BGR uint8 images.
+        generated_images: List of (H, W, 3) BGR uint8 images.
+        batch_size: Batch size for feature extraction.
+        device: "cuda" or "cpu".
+
+    Returns:
+        FID score (lower = better).
+    """
+    if not HAS_TORCH:
+        raise ImportError("PyTorch required for FID computation")
+
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    dev = torch.device(device)
+
+    model = _load_inception_v3().to(dev)
+
+    real_ds = NumpyArrayDataset(real_images)
+    gen_ds = NumpyArrayDataset(generated_images)
+
+    real_loader = DataLoader(real_ds, batch_size=batch_size, num_workers=0)
+    gen_loader = DataLoader(gen_ds, batch_size=batch_size, num_workers=0)
+
+    real_features = _extract_features(model, real_loader, dev)
+    gen_features = _extract_features(model, gen_loader, dev)
+
+    mu_real, sigma_real = _compute_statistics(real_features)
+    mu_gen, sigma_gen = _compute_statistics(gen_features)
+
+    return _calculate_fid(mu_real, sigma_real, mu_gen, sigma_gen)

landmarkdiff/hyperparam.py

@@ -0,0 +1,347 @@
+"""Hyperparameter search utilities for systematic ControlNet tuning.
+
+Supports grid search, random search, and Bayesian-inspired adaptive search
+over training hyperparameters. Generates YAML configs for each trial and
+tracks results for comparison.
+
+Usage:
+    from landmarkdiff.hyperparam import HyperparamSearch, SearchSpace
+
+    space = SearchSpace()
+    space.add_float("learning_rate", 1e-6, 1e-4, log_scale=True)
+    space.add_choice("optimizer", ["adamw", "adam8bit"])
+    space.add_int("batch_size", 2, 8, step=2)
+
+    search = HyperparamSearch(space, output_dir="hp_search")
+    for trial in search.generate_trials(strategy="random", n_trials=20):
+        print(trial.config)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+def _to_native(val: Any) -> Any:
+    """Convert numpy/non-standard types to native Python for YAML serialization."""
+    if hasattr(val, "item"):  # numpy scalar
+        return val.item()
+    return val
+
+
+@dataclass
+class ParamSpec:
+    """Specification for a single hyperparameter."""
+
+    name: str
+    param_type: str  # "float", "int", "choice"
+    low: float | None = None
+    high: float | None = None
+    step: float | None = None
+    log_scale: bool = False
+    choices: list[Any] | None = None
+
+    def sample(self, rng) -> Any:
+        """Sample a value from this parameter spec."""
+        if self.param_type == "choice":
+            return rng.choice(self.choices)
+        elif self.param_type == "float":
+            if self.log_scale:
+                log_low = math.log(self.low)
+                log_high = math.log(self.high)
+                return float(math.exp(rng.uniform(log_low, log_high)))
+            return float(rng.uniform(self.low, self.high))
+        elif self.param_type == "int":
+            if self.step and self.step > 1:
+                n_steps = int((self.high - self.low) / self.step) + 1
+                idx = rng.integers(0, n_steps)
+                return int(self.low + idx * self.step)
+            return int(rng.integers(int(self.low), int(self.high) + 1))
+        raise ValueError(f"Unknown param type: {self.param_type}")
+
+    def grid_values(self, n_points: int = 5) -> list[Any]:
+        """Generate grid values for this parameter."""
+        if self.param_type == "choice":
+            return list(self.choices)
+        elif self.param_type == "int":
+            if self.step and self.step > 1:
+                vals = []
+                v = self.low
+                while v <= self.high:
+                    vals.append(int(v))
+                    v += self.step
+                return vals
+            return list(range(int(self.low), int(self.high) + 1))
+        elif self.param_type == "float":
+            if self.log_scale:
+                log_low = math.log(self.low)
+                log_high = math.log(self.high)
+                return [
+                    float(math.exp(log_low + i * (log_high - log_low) / (n_points - 1)))
+                    for i in range(n_points)
+                ]
+            return [
+                float(self.low + i * (self.high - self.low) / (n_points - 1))
+                for i in range(n_points)
+            ]
+        return []
+
+
+class SearchSpace:
+    """Define the hyperparameter search space."""
+
+    def __init__(self) -> None:
+        self.params: dict[str, ParamSpec] = {}
+
+    def add_float(
+        self,
+        name: str,
+        low: float,
+        high: float,
+        log_scale: bool = False,
+    ) -> SearchSpace:
+        """Add a continuous float parameter."""
+        self.params[name] = ParamSpec(
+            name=name,
+            param_type="float",
+            low=low,
+            high=high,
+            log_scale=log_scale,
+        )
+        return self
+
+    def add_int(
+        self,
+        name: str,
+        low: int,
+        high: int,
+        step: int = 1,
+    ) -> SearchSpace:
+        """Add an integer parameter."""
+        self.params[name] = ParamSpec(
+            name=name,
+            param_type="int",
+            low=low,
+            high=high,
+            step=step,
+        )
+        return self
+
+    def add_choice(self, name: str, choices: list[Any]) -> SearchSpace:
+        """Add a categorical parameter."""
+        self.params[name] = ParamSpec(
+            name=name,
+            param_type="choice",
+            choices=choices,
+        )
+        return self
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.params
+
+
+@dataclass
+class Trial:
+    """A single hyperparameter trial."""
+
+    trial_id: str
+    config: dict[str, Any]
+    result: dict[str, float] = field(default_factory=dict)
+    status: str = "pending"  # pending, running, completed, failed
+
+    @property
+    def config_hash(self) -> str:
+        """Short hash of the config for deduplication."""
+        s = json.dumps(self.config, sort_keys=True, default=str)
+        return hashlib.md5(s.encode()).hexdigest()[:8]
+
+
+class HyperparamSearch:
+    """Hyperparameter search engine.
+
+    Args:
+        space: Search space definition.
+        output_dir: Directory to save trial configs and results.
+        seed: Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        space: SearchSpace,
+        output_dir: str | Path = "hp_search",
+        seed: int = 42,
+    ) -> None:
+        self.space = space
+        self.output_dir = Path(output_dir)
+        self.seed = seed
+        self.trials: list[Trial] = []
+
+    def generate_trials(
+        self,
+        strategy: str = "random",
+        n_trials: int = 20,
+        grid_points: int = 5,
+    ) -> list[Trial]:
+        """Generate trial configurations.
+
+        Args:
+            strategy: "random" or "grid".
+            n_trials: Number of trials for random search.
+            grid_points: Points per continuous dimension for grid search.
+
+        Returns:
+            List of Trial objects with configs.
+        """
+        if strategy == "grid":
+            trials = self._grid_search(grid_points)
+        elif strategy == "random":
+            trials = self._random_search(n_trials)
+        else:
+            raise ValueError(f"Unknown strategy: {strategy}. Use 'random' or 'grid'.")
+
+        self.trials.extend(trials)
+        return trials
+
+    def _random_search(self, n_trials: int) -> list[Trial]:
+        """Generate random trial configs."""
+        import numpy as np
+
+        rng = np.random.default_rng(self.seed)
+        seen_hashes: set[str] = set()
+        trials: list[Trial] = []
+
+        max_attempts = n_trials * 10
+        attempts = 0
+        while len(trials) < n_trials and attempts < max_attempts:
+            attempts += 1
+            config = {name: spec.sample(rng) for name, spec in self.space.params.items()}
+            trial = Trial(
+                trial_id=f"trial_{len(trials):04d}",
+                config=config,
+            )
+            if trial.config_hash not in seen_hashes:
+                seen_hashes.add(trial.config_hash)
+                trials.append(trial)
+
+        return trials
+
+    def _grid_search(self, grid_points: int) -> list[Trial]:
+        """Generate grid search configs."""
+        import itertools
+
+        param_names = list(self.space.params.keys())
+        param_values = [self.space.params[name].grid_values(grid_points) for name in param_names]
+
+        trials = []
+        for combo in itertools.product(*param_values):
+            config = dict(zip(param_names, combo))
+            trial = Trial(
+                trial_id=f"trial_{len(trials):04d}",
+                config=config,
+            )
+            trials.append(trial)
+
+        return trials
+
+    def record_result(
+        self,
+        trial_id: str,
+        metrics: dict[str, float],
+    ) -> None:
+        """Record results for a trial."""
+        for trial in self.trials:
+            if trial.trial_id == trial_id:
+                trial.result = metrics
+                trial.status = "completed"
+                return
+        raise KeyError(f"Trial {trial_id} not found")
+
+    def best_trial(
+        self,
+        metric: str = "loss",
+        lower_is_better: bool = True,
+    ) -> Trial | None:
+        """Get the best completed trial by a metric."""
+        completed = [t for t in self.trials if t.status == "completed" and metric in t.result]
+        if not completed:
+            return None
+        return (min if lower_is_better else max)(completed, key=lambda t: t.result[metric])
+
+    def save_configs(self) -> Path:
+        """Save all trial configs as YAML files.
+
+        Returns:
+            Output directory path.
+        """
+        import yaml
+
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        for trial in self.trials:
+            cfg_path = self.output_dir / f"{trial.trial_id}.yaml"
+            # Convert numpy types to native Python for YAML serialization
+            native_config = {k: _to_native(v) for k, v in trial.config.items()}
+            with open(cfg_path, "w") as f:
+                yaml.safe_dump(
+                    {"trial_id": trial.trial_id, **native_config},
+                    f,
+                    default_flow_style=False,
+                )
+
+        # Save summary index
+        index = {
+            "seed": self.seed,
+            "n_trials": len(self.trials),
+            "params": {
+                name: {
+                    "type": spec.param_type,
+                    "low": spec.low,
+                    "high": spec.high,
+                    "choices": spec.choices,
+                    "log_scale": spec.log_scale,
+                }
+                for name, spec in self.space.params.items()
+            },
+        }
+        with open(self.output_dir / "search_index.json", "w") as f:
+            json.dump(index, f, indent=2, default=str)
+
+        return self.output_dir
+
+    def results_table(self) -> str:
+        """Format results as a text table."""
+        completed = [t for t in self.trials if t.status == "completed"]
+        if not completed:
+            return "No completed trials."
+
+        # Collect all metric names
+        metric_names = sorted(set().union(*(t.result.keys() for t in completed)))
+        param_names = sorted(self.space.params.keys())
+
+        # Header
+        cols = ["Trial"] + param_names + metric_names
+        lines = [" | ".join(f"{c:>12s}" for c in cols)]
+        lines.append("-" * len(lines[0]))
+
+        # Rows
+        for trial in completed:
+            parts = [f"{trial.trial_id:>12s}"]
+            for p in param_names:
+                val = trial.config.get(p, "")
+                if isinstance(val, float):
+                    parts.append(f"{val:>12.6f}")
+                else:
+                    parts.append(f"{val!s:>12s}")
+            for m in metric_names:
+                val = trial.result.get(m, float("nan"))
+                parts.append(f"{val:>12.4f}")
+            lines.append(" | ".join(parts))
+
+        return "\n".join(lines)