| import os |
| from typing import Iterable, Union |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from PIL import Image |
| from torchvision import transforms |
| from torchvision.models import inception_v3 |
| from torchvision.models.inception import InceptionA, InceptionC, InceptionE |
|
|
| ImageInput = Union[str, os.PathLike, Image.Image] |
|
|
| IMAGE_EXTENSIONS = {".bmp", ".jpg", ".jpeg", ".pgm", ".png", ".ppm", ".tif", ".tiff", ".webp"} |
|
|
|
|
| def _image_files(path: Union[str, os.PathLike]): |
| path = os.fspath(path) |
| if os.path.isfile(path): |
| if os.path.splitext(path)[1].lower() not in IMAGE_EXTENSIONS: |
| raise ValueError(f"Unsupported image extension for FID: {path}") |
| return [path] |
| if not os.path.exists(path): |
| raise FileNotFoundError(f"FID path does not exist: {path}") |
| files = [] |
| for root, dirs, names in os.walk(path): |
| dirs.sort() |
| for name in sorted(names): |
| if os.path.splitext(name)[1].lower() in IMAGE_EXTENSIONS: |
| files.append(os.path.join(root, name)) |
| if not files: |
| raise ValueError(f"No images found under {path}.") |
| return files |
|
|
|
|
| class _ImageDataset(torch.utils.data.Dataset): |
| def __init__(self, images: Iterable[ImageInput], transform): |
| self.images = list(images) |
| self.transform = transform |
|
|
| def __len__(self): |
| return len(self.images) |
|
|
| def __getitem__(self, index): |
| image = self.images[index] |
| if isinstance(image, (str, os.PathLike)): |
| image = Image.open(image) |
| if not isinstance(image, Image.Image): |
| raise TypeError(f"FID expects PIL images or image paths, but received {type(image)}.") |
| return self.transform(image.convert("RGB")) |
|
|
|
|
| class FIDInceptionModel(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.model = _fid_inception_v3() |
|
|
| def forward(self, images): |
| images = 2 * images - 1 |
| return self.model(images) |
|
|
|
|
| def _fid_inception_v3(weights_path: str = None): |
| model = inception_v3(weights=None, aux_logits=False, num_classes=1008, init_weights=False) |
| model.Mixed_5b = _FIDInceptionA(192, pool_features=32) |
| model.Mixed_5c = _FIDInceptionA(256, pool_features=64) |
| model.Mixed_5d = _FIDInceptionA(288, pool_features=64) |
| model.Mixed_6b = _FIDInceptionC(768, channels_7x7=128) |
| model.Mixed_6c = _FIDInceptionC(768, channels_7x7=160) |
| model.Mixed_6d = _FIDInceptionC(768, channels_7x7=160) |
| model.Mixed_6e = _FIDInceptionC(768, channels_7x7=192) |
| model.Mixed_7b = _FIDInceptionE1(1280) |
| model.Mixed_7c = _FIDInceptionE2(2048) |
| if weights_path is not None: |
| model.load_state_dict(torch.load(weights_path, map_location="cpu")) |
| model.fc = nn.Identity() |
| return model |
|
|
|
|
| class _FIDInceptionA(InceptionA): |
| def forward(self, x): |
| branch1x1 = self.branch1x1(x) |
|
|
| branch5x5 = self.branch5x5_1(x) |
| branch5x5 = self.branch5x5_2(branch5x5) |
|
|
| branch3x3dbl = self.branch3x3dbl_1(x) |
| branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) |
| branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) |
|
|
| branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False) |
| branch_pool = self.branch_pool(branch_pool) |
|
|
| return torch.cat([branch1x1, branch5x5, branch3x3dbl, branch_pool], 1) |
|
|
|
|
| class _FIDInceptionC(InceptionC): |
| def forward(self, x): |
| branch1x1 = self.branch1x1(x) |
|
|
| branch7x7 = self.branch7x7_1(x) |
| branch7x7 = self.branch7x7_2(branch7x7) |
| branch7x7 = self.branch7x7_3(branch7x7) |
|
|
| branch7x7dbl = self.branch7x7dbl_1(x) |
| branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl) |
| branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl) |
| branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl) |
| branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl) |
|
|
| branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False) |
| branch_pool = self.branch_pool(branch_pool) |
|
|
| return torch.cat([branch1x1, branch7x7, branch7x7dbl, branch_pool], 1) |
|
|
|
|
| class _FIDInceptionE1(InceptionE): |
| def forward(self, x): |
| branch1x1 = self.branch1x1(x) |
|
|
| branch3x3 = self.branch3x3_1(x) |
| branch3x3 = torch.cat([self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)], 1) |
|
|
| branch3x3dbl = self.branch3x3dbl_1(x) |
| branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) |
| branch3x3dbl = torch.cat([self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl)], 1) |
|
|
| branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False) |
| branch_pool = self.branch_pool(branch_pool) |
|
|
| return torch.cat([branch1x1, branch3x3, branch3x3dbl, branch_pool], 1) |
|
|
|
|
| class _FIDInceptionE2(InceptionE): |
| def forward(self, x): |
| branch1x1 = self.branch1x1(x) |
|
|
| branch3x3 = self.branch3x3_1(x) |
| branch3x3 = torch.cat([self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)], 1) |
|
|
| branch3x3dbl = self.branch3x3dbl_1(x) |
| branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) |
| branch3x3dbl = torch.cat([self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl)], 1) |
|
|
| branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1) |
| branch_pool = self.branch_pool(branch_pool) |
|
|
| return torch.cat([branch1x1, branch3x3, branch3x3dbl, branch_pool], 1) |
|
|
|
|
| class FIDModel(torch.nn.Module): |
| def __init__(self, model: torch.nn.Module, device: Union[str, torch.device] = "cpu", batch_size: int = 50, num_workers: int = 0): |
| super().__init__() |
| self.model = model |
| self.batch_size = batch_size |
| self.num_workers = num_workers |
| self.transform = transforms.Compose( |
| [ |
| transforms.Resize((299, 299), interpolation=transforms.InterpolationMode.BICUBIC), |
| transforms.ToTensor(), |
| ] |
| ) |
| self.to(device) |
|
|
| @property |
| def device(self): |
| try: |
| return next(self.model.parameters()).device |
| except StopIteration: |
| return torch.device("cpu") |
|
|
| def _as_images(self, images): |
| if isinstance(images, (str, os.PathLike)): |
| return _image_files(images) |
| if isinstance(images, Image.Image): |
| return [images] |
| return list(images) |
|
|
| @torch.no_grad() |
| def get_activations(self, images, batch_size: int = None, num_workers: int = None): |
| images = self._as_images(images) |
| batch_size = self.batch_size if batch_size is None else batch_size |
| num_workers = self.num_workers if num_workers is None else num_workers |
| dataset = _ImageDataset(images, transform=self.transform) |
| dataloader = torch.utils.data.DataLoader(dataset, batch_size=min(batch_size, len(dataset)), shuffle=False, num_workers=num_workers) |
| activations = [] |
| self.model.eval() |
| for batch in dataloader: |
| batch = batch.to(self.device) |
| features = self.model(batch) |
| if isinstance(features, tuple): |
| features = features[0] |
| if features.ndim == 4: |
| features = F.adaptive_avg_pool2d(features, output_size=(1, 1)).flatten(1) |
| activations.append(features.detach().cpu().to(torch.float64)) |
| return torch.cat(activations, dim=0) |
|
|
| def statistics(self, images, batch_size: int = None, num_workers: int = None): |
| activations = self.get_activations(images, batch_size=batch_size, num_workers=num_workers) |
| return self.activation_statistics(activations) |
|
|
| @staticmethod |
| def activation_statistics(activations): |
| activations = activations.to(torch.float64) |
| mean = activations.mean(dim=0) |
| centered = activations - mean |
| if activations.shape[0] <= 1: |
| covariance = torch.zeros((activations.shape[1], activations.shape[1]), dtype=torch.float64) |
| else: |
| covariance = centered.T @ centered / (activations.shape[0] - 1) |
| return mean, covariance |
|
|
| @staticmethod |
| def _sqrtm_psd(matrix, eps: float = 1e-10): |
| matrix = (matrix + matrix.T) * 0.5 |
| eigenvalues, eigenvectors = torch.linalg.eigh(matrix) |
| eigenvalues = eigenvalues.clamp_min(eps).sqrt() |
| return (eigenvectors * eigenvalues.unsqueeze(0)) @ eigenvectors.T |
|
|
| @classmethod |
| def frechet_distance(cls, mean1, covariance1, mean2, covariance2, eps: float = 1e-6): |
| mean1 = mean1.to(torch.float64) |
| covariance1 = covariance1.to(torch.float64) |
| mean2 = mean2.to(torch.float64) |
| covariance2 = covariance2.to(torch.float64) |
| diff = mean1 - mean2 |
| offset = torch.eye(covariance1.shape[0], dtype=torch.float64) * eps |
| sqrt_cov1 = cls._sqrtm_psd(covariance1 + offset) |
| covmean = cls._sqrtm_psd(sqrt_cov1 @ (covariance2 + offset) @ sqrt_cov1) |
| distance = diff.dot(diff) + torch.trace(covariance1) + torch.trace(covariance2) - 2 * torch.trace(covmean) |
| return distance.clamp_min(0) |
|
|
| def compute(self, reference_images, generated_images, batch_size: int = None, num_workers: int = None): |
| mean1, covariance1 = self.statistics(reference_images, batch_size=batch_size, num_workers=num_workers) |
| mean2, covariance2 = self.statistics(generated_images, batch_size=batch_size, num_workers=num_workers) |
| return self.frechet_distance(mean1, covariance1, mean2, covariance2) |
|
|
| def forward(self, reference_images, generated_images): |
| return self.compute(reference_images, generated_images) |
|
|
