cache-dit 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

cache_dit/metrics/fid.py

@@ -0,0 +1,409 @@
+import os
+import cv2
+import pathlib
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from scipy import linalg
+import torch
+import torchvision.transforms as TF
+from torch.nn.functional import adaptive_avg_pool2d
+from cache_dit.metrics.inception import InceptionV3
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ImagePathDataset(torch.utils.data.Dataset):
+    def __init__(self, files_or_imgs, transforms=None):
+        self.files_or_imgs = files_or_imgs
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.files_or_imgs)
+
+    def __getitem__(self, i):
+        file_or_img = self.files_or_imgs[i]
+        if isinstance(file_or_img, (str, pathlib.Path)):
+            img = Image.open(file_or_img).convert("RGB")
+        elif isinstance(file_or_img, np.ndarray):
+            # Assume the img is a standard OpenCV image.
+            img = cv2.cvtColor(file_or_img, cv2.COLOR_BGR2RGB)
+            img = Image.fromarray(img)
+        else:
+            raise ValueError(
+                "file_or_img must be a file path or an OpenCV image."
+            )
+        if self.transforms is not None:
+            img = self.transforms(img)
+        return img
+
+
+def get_activations(
+    files_or_imgs,
+    model,
+    batch_size=50,
+    dims=2048,
+    device="cpu",
+    num_workers=1,
+    disable_tqdm=True,
+):
+    """Calculates the activations of the pool_3 layer for all images.
+
+    Params:
+    -- files_or_imgs : List of image files paths or OpenCV image
+    -- model         : Instance of inception model
+    -- batch_size    : Batch size of images for the model to process at once.
+                       Make sure that the number of samples is a multiple of
+                       the batch size, otherwise some samples are ignored. This
+                       behavior is retained to match the original FID score
+                       implementation.
+    -- dims          : Dimensionality of features returned by Inception
+    -- device        : Device to run calculations
+    -- num_workers   : Number of parallel dataloader workers
+
+    Returns:
+    -- A numpy array of dimension (num images, dims) that contains the
+       activations of the given tensor when feeding inception with the
+       query tensor.
+    """
+    model.eval()
+
+    if batch_size > len(files_or_imgs):
+        logger.info(
+            (
+                "Warning: batch size is bigger than the data size. "
+                "Setting batch size to data size"
+            )
+        )
+        batch_size = len(files_or_imgs)
+
+    dataset = ImagePathDataset(files_or_imgs, transforms=TF.ToTensor())
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=num_workers,
+    )
+
+    pred_arr = np.empty((len(files_or_imgs), dims))
+
+    start_idx = 0
+
+    for batch in tqdm(dataloader, disable=disable_tqdm):
+        batch = batch.to(device)
+
+        with torch.no_grad():
+            pred = model(batch)[0]
+
+        # If model output is not scalar, apply global spatial average pooling.
+        # This happens if you choose a dimensionality not equal 2048.
+        if pred.size(2) != 1 or pred.size(3) != 1:
+            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+
+        pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+
+        pred_arr[start_idx : start_idx + pred.shape[0]] = pred
+
+        start_idx = start_idx + pred.shape[0]
+
+    return pred_arr
+
+
+def calculate_frechet_distance(
+    mu1,
+    sigma1,
+    mu2,
+    sigma2,
+    eps=1e-6,
+):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert (
+        mu1.shape == mu2.shape
+    ), "Training and test mean vectors have different lengths"
+    assert (
+        sigma1.shape == sigma2.shape
+    ), "Training and test covariances have different dimensions"
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = (
+            "fid calculation produces singular product; "
+            "adding %s to diagonal of cov estimates"
+        ) % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError("Imaginary component {}".format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+
+
+def calculate_activation_statistics(
+    files_or_imgs,
+    model,
+    batch_size=50,
+    dims=2048,
+    device="cpu",
+    num_workers=1,
+    disable_tqdm=True,
+):
+    """Calculation of the statistics used by the FID.
+    Params:
+    -- files_or_imgs : List of image files paths or OpenCV image
+    -- model         : Instance of inception model
+    -- batch_size    : Batch size of images for the model to process at once.
+                       Make sure that the number of samples is a multiple of
+                       the batch size, otherwise some samples are ignored. This
+                       behavior is retained to match the original FID score
+                       implementation.
+    -- dims          : Dimensionality of features returned by Inception
+    -- device        : Device to run calculations
+    -- num_workers   : Number of parallel dataloader workers
+
+    Returns:
+    -- mu    : The mean over samples of the activations of the pool_3 layer of
+               the inception model.
+    -- sigma : The covariance matrix of the activations of the pool_3 layer of
+               the inception model.
+    """
+    act = get_activations(
+        files_or_imgs,
+        model,
+        batch_size,
+        dims,
+        device,
+        num_workers,
+        disable_tqdm,
+    )
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+_IMAGE_EXTENSIONS = {
+    "bmp",
+    "jpg",
+    "jpeg",
+    "pgm",
+    "png",
+    "ppm",
+    "tif",
+    "tiff",
+    "webp",
+}
+
+
+class FrechetInceptionDistance:
+    IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS
+
+    def __init__(
+        self,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        dims: int = 2048,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        disable_tqdm: bool = True,
+    ):
+        # https://github.com/mseitzer/pytorch-fid/src/pytorch_fid/fid_score.py
+        self.dims = dims
+        self.device = device
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        self.disable_tqdm = disable_tqdm
+        self.block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims]
+        self.model = InceptionV3([self.block_idx]).to(self.device)
+        self.model = self.model.eval()
+
+    def compute_fid(
+        self,
+        image_true: np.ndarray | str,
+        image_test: np.ndarray | str,
+    ):
+        """Calculates the FID of two file paths
+        FID = FrechetInceptionDistance()
+        img_fid = FID.compute_fid("img_true.png", "img_test.png")
+        img_dir_fid = FID.compute_fid("img_true_dir", "img_test_dir")
+        """
+        if isinstance(image_true, str) or isinstance(image_test, str):
+            if os.path.isfile(image_true) or os.path.isfile(image_test):
+                assert os.path.exists(image_true)
+                assert os.path.exists(image_test)
+                assert image_true.split(".")[-1] in self.IMAGE_EXTENSIONS
+                assert image_test.split(".")[-1] in self.IMAGE_EXTENSIONS
+                image_true_files = [image_true]
+                image_test_files = [image_test]
+            else:
+                # glob image files from dir
+                assert os.path.isdir(image_true)
+                assert os.path.isdir(image_test)
+                image_true_dir = pathlib.Path(image_true)
+                image_true_files = sorted(
+                    [
+                        file
+                        for ext in self.IMAGE_EXTENSIONS
+                        for file in image_true_dir.rglob("*.{}".format(ext))
+                    ]
+                )
+                image_test_dir = pathlib.Path(image_test)
+                image_test_files = sorted(
+                    [
+                        file
+                        for ext in self.IMAGE_EXTENSIONS
+                        for file in image_test_dir.rglob("*.{}".format(ext))
+                    ]
+                )
+                image_true_files = [
+                    file.as_posix() for file in image_true_files
+                ]
+                image_test_files = [
+                    file.as_posix() for file in image_test_files
+                ]
+                logger.debug(f"image_true_files: {image_true_files}")
+                logger.debug(f"image_test_files: {image_test_files}")
+                assert len(image_true_files) == len(image_test_files)
+                for image_true, image_test in zip(
+                    image_true_files, image_test_files
+                ):
+                    assert os.path.basename(image_true) == os.path.basename(
+                        image_test
+                    ), f"image_true:{image_true} != image_test: {image_test}"
+        else:
+            image_true_files = [image_true]
+            image_test_files = [image_test]
+
+        batch_size = min(16, self.batch_size)
+        batch_size = min(batch_size, len(image_test_files))
+        m1, s1 = calculate_activation_statistics(
+            image_true_files,
+            self.model,
+            batch_size,
+            self.dims,
+            self.device,
+            self.num_workers,
+            self.disable_tqdm,
+        )
+        m2, s2 = calculate_activation_statistics(
+            image_test_files,
+            self.model,
+            batch_size,
+            self.dims,
+            self.device,
+            self.num_workers,
+            self.disable_tqdm,
+        )
+        fid_value = calculate_frechet_distance(
+            m1,
+            s1,
+            m2,
+            s2,
+        )
+
+        return fid_value, len(image_true_files)
+
+    def compute_video_fid(
+        self,
+        video_true: str,
+        video_test: str,
+    ):
+        cap1 = cv2.VideoCapture(video_true)
+        cap2 = cv2.VideoCapture(video_test)
+
+        if not cap1.isOpened() or not cap2.isOpened():
+            logger.error("Could not open video files")
+            return None, None
+
+        frame_count = min(
+            int(cap1.get(cv2.CAP_PROP_FRAME_COUNT)),
+            int(cap2.get(cv2.CAP_PROP_FRAME_COUNT)),
+        )
+
+        valid_frames = 0
+        video_true_frames = []
+        video_test_frames = []
+
+        logger.debug(f"Total frames: {frame_count}")
+
+        while True:
+            ret1, frame1 = cap1.read()
+            ret2, frame2 = cap2.read()
+
+            if not ret1 or not ret2:
+                break
+
+            video_true_frames.append(frame1)
+            video_test_frames.append(frame2)
+
+            valid_frames += 1
+
+        cap1.release()
+        cap2.release()
+
+        if valid_frames <= 0:
+            return None, None
+
+        batch_size = min(16, self.batch_size)
+        m1, s1 = calculate_activation_statistics(
+            video_true_frames,
+            self.model,
+            batch_size,
+            self.dims,
+            self.device,
+            self.num_workers,
+            self.disable_tqdm,
+        )
+        m2, s2 = calculate_activation_statistics(
+            video_test_frames,
+            self.model,
+            batch_size,
+            self.dims,
+            self.device,
+            self.num_workers,
+            self.disable_tqdm,
+        )
+        fid_value = calculate_frechet_distance(
+            m1,
+            s1,
+            m2,
+            s2,
+        )
+
+        return fid_value, valid_frames