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

cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.5'
-__version_tuple__ = version_tuple = (0, 2, 5)
+__version__ = version = '0.2.6'
+__version_tuple__ = version_tuple = (0, 2, 6)

cache_dit/metrics/__init__.py

@@ -0,0 +1,12 @@
+from cache_dit.metrics.metrics import compute_psnr
+from cache_dit.metrics.metrics import compute_ssim
+from cache_dit.metrics.metrics import compute_mse
+from cache_dit.metrics.metrics import compute_video_psnr
+from cache_dit.metrics.metrics import compute_video_ssim
+from cache_dit.metrics.metrics import compute_video_mse
+from cache_dit.metrics.metrics import entrypoint
+from cache_dit.metrics.fid import FrechetInceptionDistance
+
+
+def main():
+    entrypoint()

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

cache_dit/metrics/inception.py

@@ -0,0 +1,353 @@
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+# Adapted from: https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/inception.py
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# Inception weights ported to Pytorch from
+# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+FID_WEIGHTS_URL = "https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth"  # noqa: E501
+
+
+class InceptionV3(nn.Module):
+    """Pretrained InceptionV3 network returning feature maps"""
+
+    # Index of default block of inception to return,
+    # corresponds to output of final average pooling
+    DEFAULT_BLOCK_INDEX = 3
+
+    # Maps feature dimensionality to their output blocks indices
+    BLOCK_INDEX_BY_DIM = {
+        64: 0,  # First max pooling features
+        192: 1,  # Second max pooling featurs
+        768: 2,  # Pre-aux classifier features
+        2048: 3,  # Final average pooling features
+    }
+
+    def __init__(
+        self,
+        output_blocks=(DEFAULT_BLOCK_INDEX,),
+        resize_input=True,
+        normalize_input=True,
+        requires_grad=False,
+        use_fid_inception=True,
+    ):
+        """Build pretrained InceptionV3
+
+        Parameters
+        ----------
+        output_blocks : list of int
+            Indices of blocks to return features of. Possible values are:
+                - 0: corresponds to output of first max pooling
+                - 1: corresponds to output of second max pooling
+                - 2: corresponds to output which is fed to aux classifier
+                - 3: corresponds to output of final average pooling
+        resize_input : bool
+            If true, bilinearly resizes input to width and height 299 before
+            feeding input to model. As the network without fully connected
+            layers is fully convolutional, it should be able to handle inputs
+            of arbitrary size, so resizing might not be strictly needed
+        normalize_input : bool
+            If true, scales the input from range (0, 1) to the range the
+            pretrained Inception network expects, namely (-1, 1)
+        requires_grad : bool
+            If true, parameters of the model require gradients. Possibly useful
+            for finetuning the network
+        use_fid_inception : bool
+            If true, uses the pretrained Inception model used in Tensorflow's
+            FID implementation. If false, uses the pretrained Inception model
+            available in torchvision. The FID Inception model has different
+            weights and a slightly different structure from torchvision's
+            Inception model. If you want to compute FID scores, you are
+            strongly advised to set this parameter to true to get comparable
+            results.
+        """
+        super(InceptionV3, self).__init__()
+
+        self.resize_input = resize_input
+        self.normalize_input = normalize_input
+        self.output_blocks = sorted(output_blocks)
+        self.last_needed_block = max(output_blocks)
+
+        assert (
+            self.last_needed_block <= 3
+        ), "Last possible output block index is 3"
+
+        self.blocks = nn.ModuleList()
+
+        if use_fid_inception:
+            inception = fid_inception_v3()
+        else:
+            inception = _inception_v3(weights="DEFAULT")
+
+        # Block 0: input to maxpool1
+        block0 = [
+            inception.Conv2d_1a_3x3,
+            inception.Conv2d_2a_3x3,
+            inception.Conv2d_2b_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2),
+        ]
+        self.blocks.append(nn.Sequential(*block0))
+
+        # Block 1: maxpool1 to maxpool2
+        if self.last_needed_block >= 1:
+            block1 = [
+                inception.Conv2d_3b_1x1,
+                inception.Conv2d_4a_3x3,
+                nn.MaxPool2d(kernel_size=3, stride=2),
+            ]
+            self.blocks.append(nn.Sequential(*block1))
+
+        # Block 2: maxpool2 to aux classifier
+        if self.last_needed_block >= 2:
+            block2 = [
+                inception.Mixed_5b,
+                inception.Mixed_5c,
+                inception.Mixed_5d,
+                inception.Mixed_6a,
+                inception.Mixed_6b,
+                inception.Mixed_6c,
+                inception.Mixed_6d,
+                inception.Mixed_6e,
+            ]
+            self.blocks.append(nn.Sequential(*block2))
+
+        # Block 3: aux classifier to final avgpool
+        if self.last_needed_block >= 3:
+            block3 = [
+                inception.Mixed_7a,
+                inception.Mixed_7b,
+                inception.Mixed_7c,
+                nn.AdaptiveAvgPool2d(output_size=(1, 1)),
+            ]
+            self.blocks.append(nn.Sequential(*block3))
+
+        for param in self.parameters():
+            param.requires_grad = requires_grad
+
+    def forward(self, inp):
+        """Get Inception feature maps
+
+        Parameters
+        ----------
+        inp : torch.autograd.Variable
+            Input tensor of shape Bx3xHxW. Values are expected to be in
+            range (0, 1)
+
+        Returns
+        -------
+        List of torch.autograd.Variable, corresponding to the selected output
+        block, sorted ascending by index
+        """
+        outp = []
+        x = inp
+
+        if self.resize_input:
+            x = F.interpolate(
+                x, size=(299, 299), mode="bilinear", align_corners=False
+            )
+
+        if self.normalize_input:
+            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x)
+            if idx in self.output_blocks:
+                outp.append(x)
+
+            if idx == self.last_needed_block:
+                break
+
+        return outp
+
+
+def _inception_v3(*args, **kwargs):
+    """Wraps `torchvision.models.inception_v3`"""
+    try:
+        version = tuple(map(int, torchvision.__version__.split(".")[:2]))
+    except ValueError:
+        # Just a caution against weird version strings
+        version = (0,)
+
+    # Skips default weight inititialization if supported by torchvision
+    # version. See https://github.com/mseitzer/pytorch-fid/issues/28.
+    if version >= (0, 6):
+        kwargs["init_weights"] = False
+
+    # Backwards compatibility: `weights` argument was handled by `pretrained`
+    # argument prior to version 0.13.
+    if version < (0, 13) and "weights" in kwargs:
+        if kwargs["weights"] == "DEFAULT":
+            kwargs["pretrained"] = True
+        elif kwargs["weights"] is None:
+            kwargs["pretrained"] = False
+        else:
+            raise ValueError(
+                "weights=={} not supported in torchvision {}".format(
+                    kwargs["weights"], torchvision.__version__
+                )
+            )
+        del kwargs["weights"]
+
+    return torchvision.models.inception_v3(*args, **kwargs)
+
+
+def fid_inception_v3():
+    """Build pretrained Inception model for FID computation
+
+    The Inception model for FID computation uses a different set of weights
+    and has a slightly different structure than torchvision's Inception.
+
+    This method first constructs torchvision's Inception and then patches the
+    necessary parts that are different in the FID Inception model.
+    """
+    inception = _inception_v3(num_classes=1008, aux_logits=False, weights=None)
+    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
+    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
+    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
+    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
+    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
+    inception.Mixed_7b = FIDInceptionE_1(1280)
+    inception.Mixed_7c = FIDInceptionE_2(2048)
+
+    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=False)
+    inception.load_state_dict(state_dict)
+    return inception
+
+
+class FIDInceptionA(torchvision.models.inception.InceptionA):
+    """InceptionA block patched for FID computation"""
+
+    def __init__(self, in_channels, pool_features):
+        super(FIDInceptionA, self).__init__(in_channels, pool_features)
+
+    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)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(
+            x, kernel_size=3, stride=1, padding=1, count_include_pad=False
+        )
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionC(torchvision.models.inception.InceptionC):
+    """InceptionC block patched for FID computation"""
+
+    def __init__(self, in_channels, channels_7x7):
+        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
+
+    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)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(
+            x, kernel_size=3, stride=1, padding=1, count_include_pad=False
+        )
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_1(torchvision.models.inception.InceptionE):
+    """First InceptionE block patched for FID computation"""
+
+    def __init__(self, in_channels):
+        super(FIDInceptionE_1, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(
+            x, kernel_size=3, stride=1, padding=1, count_include_pad=False
+        )
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_2(torchvision.models.inception.InceptionE):
+    """Second InceptionE block patched for FID computation"""
+
+    def __init__(self, in_channels):
+        super(FIDInceptionE_2, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: The FID Inception model uses max pooling instead of average
+        # pooling. This is likely an error in this specific Inception
+        # implementation, as other Inception models use average pooling here
+        # (which matches the description in the paper).
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)

cache_dit/metrics/metrics.py

@@ -0,0 +1,356 @@
+import os
+import cv2
+import pathlib
+import argparse
+import numpy as np
+from functools import partial
+from skimage.metrics import mean_squared_error
+from skimage.metrics import peak_signal_noise_ratio
+from skimage.metrics import structural_similarity
+from cache_dit.metrics.fid import FrechetInceptionDistance
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def compute_psnr_file(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> float:
+    """
+    img_true = cv2.imread(img_true_file)
+    img_test = cv2.imread(img_test_file)
+    PSNR = compute_psnr(img_true, img_test)
+    """
+    if isinstance(image_true, str):
+        image_true = cv2.imread(image_true)
+    if isinstance(image_test, str):
+        image_test = cv2.imread(image_test)
+    return peak_signal_noise_ratio(
+        image_true,
+        image_test,
+    )
+
+
+def compute_mse_file(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> float:
+    """
+    img_true = cv2.imread(img_true_file)
+    img_test = cv2.imread(img_test_file)
+    MSE = compute_mse(img_true, img_test)
+    """
+    if isinstance(image_true, str):
+        image_true = cv2.imread(image_true)
+    if isinstance(image_test, str):
+        image_test = cv2.imread(image_test)
+    return mean_squared_error(
+        image_true,
+        image_test,
+    )
+
+
+def compute_ssim_file(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> float:
+    """
+    img_true = cv2.imread(img_true_file)
+    img_test = cv2.imread(img_test_file)
+    SSIM = compute_ssim(img_true, img_test)
+    """
+    if isinstance(image_true, str):
+        image_true = cv2.imread(image_true)
+    if isinstance(image_test, str):
+        image_test = cv2.imread(image_test)
+    return structural_similarity(
+        image_true,
+        image_test,
+        multichannel=True,
+        channel_axis=2,
+    )
+
+
+_IMAGE_EXTENSIONS = {
+    "bmp",
+    "jpg",
+    "jpeg",
+    "pgm",
+    "png",
+    "ppm",
+    "tif",
+    "tiff",
+    "webp",
+}
+
+
+def compute_dir_metric(
+    image_true_dir: np.ndarray | str,
+    image_test_dir: np.ndarray | str,
+    compute_file_func: callable = compute_psnr_file,
+) -> float:
+    # Image
+    if isinstance(image_true_dir, np.ndarray) or isinstance(
+        image_test_dir, np.ndarray
+    ):
+        return compute_file_func(image_true_dir, image_test_dir), 1
+    # File
+    if not os.path.isdir(image_true_dir) or not os.path.isdir(image_test_dir):
+        return compute_file_func(image_true_dir, image_test_dir), 1
+    # Dir
+    image_true_dir: pathlib.Path = pathlib.Path(image_true_dir)
+    image_true_files = sorted(
+        [
+            file
+            for ext in _IMAGE_EXTENSIONS
+            for file in image_true_dir.rglob("*.{}".format(ext))
+        ]
+    )
+    image_test_dir: pathlib.Path = pathlib.Path(image_test_dir)
+    image_test_files = sorted(
+        [
+            file
+            for ext in _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}"
+
+    total_metric = 0.0
+    valid_files = 0
+    for image_true, image_test in zip(image_true_files, image_test_files):
+        metric = compute_file_func(image_true, image_test)
+        if metric != float("inf"):
+            total_metric += metric
+            valid_files += 1
+
+    if valid_files > 0:
+        average_metric = total_metric / valid_files
+        logger.debug(f"Average: {average_metric:.2f}")
+        return average_metric, valid_files
+    else:
+        logger.debug("No valid files to compare")
+        return None, None
+
+
+def compute_video_metric(
+    video_true: str,
+    video_test: str,
+    compute_frame_func: callable = compute_psnr_file,
+) -> float:
+    """
+    video_true = "video_true.mp4"
+    video_test = "video_test.mp4"
+    PSNR = compute_video_psnr(video_true, video_test)
+    """
+    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
+
+    frame_count = min(
+        int(cap1.get(cv2.CAP_PROP_FRAME_COUNT)),
+        int(cap2.get(cv2.CAP_PROP_FRAME_COUNT)),
+    )
+
+    total_metric = 0.0
+    valid_frames = 0
+
+    logger.debug(f"Total frames: {frame_count}")
+
+    while True:
+        ret1, frame1 = cap1.read()
+        ret2, frame2 = cap2.read()
+
+        if not ret1 or not ret2:
+            break
+
+        metric = compute_frame_func(frame1, frame2)
+
+        if metric != float("inf"):
+            total_metric += metric
+            valid_frames += 1
+
+        if valid_frames % 10 == 0:
+            logger.debug(f"Processed {valid_frames}/{frame_count} frames")
+
+    cap1.release()
+    cap2.release()
+
+    if valid_frames > 0:
+        average_metric = total_metric / valid_frames
+        logger.debug(f"Average: {average_metric:.2f}")
+        return average_metric, valid_frames
+    else:
+        logger.debug("No valid frames to compare")
+        return None, None
+
+
+compute_psnr = partial(
+    compute_dir_metric,
+    compute_file_func=compute_psnr_file,
+)
+
+compute_ssim = partial(
+    compute_dir_metric,
+    compute_file_func=compute_ssim_file,
+)
+
+compute_mse = partial(
+    compute_dir_metric,
+    compute_file_func=compute_mse_file,
+)
+
+compute_video_psnr = partial(
+    compute_video_metric,
+    compute_frame_func=compute_psnr_file,
+)
+compute_video_ssim = partial(
+    compute_video_metric,
+    compute_frame_func=compute_ssim_file,
+)
+compute_video_mse = partial(
+    compute_video_metric,
+    compute_frame_func=compute_mse_file,
+)
+
+
+# Entrypoints
+def get_args():
+    parser = argparse.ArgumentParser(
+        description="CacheDiT's Metrics CLI",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    METRICS_CHOICES = [
+        "psnr",
+        "ssim",
+        "mse",
+        "fid",
+        "all",
+    ]
+    parser.add_argument(
+        "metric",
+        type=str,
+        default="psnr",
+        choices=METRICS_CHOICES,
+        help=f"Metric choices: {METRICS_CHOICES}",
+    )
+    parser.add_argument(
+        "--img-true",
+        "-i1",
+        type=str,
+        default=None,
+        help="Path to ground truth image or Dir to ground truth images",
+    )
+    parser.add_argument(
+        "--img-test",
+        "-i2",
+        type=str,
+        default=None,
+        help="Path to predicted image or Dir to predicted images",
+    )
+    parser.add_argument(
+        "--video-true",
+        "-v1",
+        type=str,
+        default=None,
+        help="Path to ground truth video",
+    )
+    parser.add_argument(
+        "--video-test",
+        "-v2",
+        type=str,
+        default=None,
+        help="Path to predicted video",
+    )
+    return parser.parse_args()
+
+
+def entrypoint():
+    args = get_args()
+    logger.debug(args)
+
+    if args.img_true is not None and args.img_test is not None:
+        if any(
+            (
+                not os.path.exists(args.img_true),
+                not os.path.exists(args.img_test),
+            )
+        ):
+            return
+        # img_true and img_test can be files or dirs
+        if args.metric == "psnr" or args.metric == "all":
+            img_psnr, n = compute_psnr(args.img_true, args.img_test)
+            logger.info(
+                f"{args.img_true} vs {args.img_test}, Num: {n}, PSNR: {img_psnr}"
+            )
+        if args.metric == "ssim" or args.metric == "all":
+            img_ssim, n = compute_ssim(args.img_true, args.img_test)
+            logger.info(
+                f"{args.img_true} vs {args.img_test}, Num: {n}, SSIM: {img_ssim}"
+            )
+        if args.metric == "mse" or args.metric == "all":
+            img_mse, n = compute_mse(args.img_true, args.img_test)
+            logger.info(
+                f"{args.img_true} vs {args.img_test}, Num: {n},  MSE: {img_mse}"
+            )
+        if args.metric == "fid" or args.metric == "all":
+            FID = FrechetInceptionDistance()
+            img_fid, n = FID.compute_fid(args.img_true, args.img_test)
+            logger.info(
+                f"{args.img_true} vs {args.img_test}, Num: {n},  FID: {img_fid}"
+            )
+    if args.video_true is not None and args.video_test is not None:
+        if any(
+            (
+                not os.path.exists(args.video_true),
+                not os.path.exists(args.video_test),
+            )
+        ):
+            return
+        if args.metric == "psnr" or args.metric == "all":
+            assert not os.path.isdir(args.video_true)
+            assert not os.path.isdir(args.video_test)
+            video_psnr, n = compute_video_psnr(args.video_true, args.video_test)
+            logger.info(
+                f"{args.video_true} vs {args.video_test}, Num: {n}, PSNR: {video_psnr}"
+            )
+        if args.metric == "ssim" or args.metric == "all":
+            assert not os.path.isdir(args.video_true)
+            assert not os.path.isdir(args.video_test)
+            video_ssim, n = compute_video_ssim(args.video_true, args.video_test)
+            logger.info(
+                f"{args.video_true} vs {args.video_test}, Num: {n}, SSIM: {video_ssim}"
+            )
+        if args.metric == "mse" or args.metric == "all":
+            assert not os.path.isdir(args.video_true)
+            assert not os.path.isdir(args.video_test)
+            video_mse, n = compute_video_mse(args.video_true, args.video_test)
+            logger.info(
+                f"{args.video_true} vs {args.video_test}, Num: {n},  MSE: {video_mse}"
+            )
+        if args.metric == "fid" or args.metric == "all":
+            assert not os.path.isdir(args.video_true)
+            assert not os.path.isdir(args.video_test)
+            FID = FrechetInceptionDistance()
+            video_fid, n = FID.compute_video_fid(
+                args.video_true, args.video_test
+            )
+            logger.info(
+                f"{args.video_true} vs {args.video_test}, Num: {n},  FID: {video_fid}"
+            )
+
+
+if __name__ == "__main__":
+    entrypoint()

{cache_dit-0.2.5.dist-info → cache_dit-0.2.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.5
+Version: 0.2.6
 Summary: 🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -13,6 +13,8 @@ Requires-Dist: packaging
 Requires-Dist: torch>=2.5.1
 Requires-Dist: transformers>=4.51.3
 Requires-Dist: diffusers>=0.33.1
+Requires-Dist: scikit-image
+Requires-Dist: scipy
 Provides-Extra: all
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
@@ -28,6 +30,8 @@ Requires-Dist: protobuf; extra == "dev"
 Requires-Dist: sentencepiece; extra == "dev"
 Requires-Dist: opencv-python-headless; extra == "dev"
 Requires-Dist: ftfy; extra == "dev"
+Requires-Dist: scikit-image; extra == "dev"
+Requires-Dist: pytorch-fid; extra == "dev"
 Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
@@ -159,6 +163,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 - [⚡️Dynamic Block Prune](#dbprune)
 - [🎉Context Parallelism](#context-parallelism)  
 - [🔥Torch Compile](#compile)
+- [⚙️Metrics CLI](#metrics)
 - [👋Contribute](#contribute)
 - [©️License](#license)
 
@@ -500,6 +505,39 @@ torch._dynamo.config.accumulated_recompile_limit = 2048  # default is 256
 
 Please check [bench.py](./bench/bench.py) for more details.
 
+
+## ⚙️Metrics CLI
+
+<div id="metrics"></div>    
+
+You can utilize the APIs provided by CacheDiT to quickly evaluate the accuracy losses caused by different cache configurations. For example:
+
+```python
+from cache_dit.metrics import compute_psnr
+from cache_dit.metrics import compute_video_psnr
+from cache_dit.metrics import FrechetInceptionDistance  # FID
+
+FID = FrechetInceptionDistance()
+image_psnr, n = compute_psnr("true.png", "test.png") # Num: n
+image_fid,  n = FID.compute_fid("true.png", "test.png")
+video_psnr, n = compute_video_psnr("true.mp4", "test.mp4")
+```
+
+Please check [test_metrics.py](./tests/test_metrics.py) for more details. Or, you can use `cache-dit-metrics-cli` tool. For examples: 
+
+```bash
+cache-dit-metrics-cli -h  # show usage
+cache-dit-metrics-cli all -v1 true.mp4 -v2 test.mp4  # compare video
+cache-dit-metrics-cli all -i1 true.png -i2 test.png  # compare image
+cache-dit-metrics-cli all -i1 true_dir -i2 test_dir  # compare image dir
+cache-dit-metrics-cli all -i1 BASELINE -i2 OPTIMIZED # compare image dir
+
+INFO 07-09 20:59:40 [metrics.py:295] BASELINE vs OPTIMIZED, Num: 1000, PSNR: 38.742413478199005
+INFO 07-09 21:00:32 [metrics.py:300] BASELINE vs OPTIMIZED, Num: 1000, SSIM: 0.9863484896791567
+INFO 07-09 21:00:45 [metrics.py:305] BASELINE vs OPTIMIZED, Num: 1000,  MSE: 12.287594770695606
+INFO 07-09 21:01:04 [metrics.py:311] BASELINE vs OPTIMIZED, Num: 1000,  FID: 5.983550108647762
+```
+
 ## 👋Contribute 
 <div id="contribute"></div>
 

{cache_dit-0.2.5.dist-info → cache_dit-0.2.6.dist-info}/RECORD

@@ -1,5 +1,5 @@
 cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=N3oBwJUFmS-AwCjqOcSlRW4GvSq-uJJMaBvoGfv1-hM,511
+cache_dit/_version.py,sha256=nObnONsicQ3YX6SG5MVBxmIp5dmRacXDauSqZijWQbY,511
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
 cache_dit/cache_factory/__init__.py,sha256=5RNuhWakvvqrOV4vkqrEBA7d-V1LwcNSsjtW14mkqK8,5255
@@ -33,8 +33,13 @@ cache_dit/compile/__init__.py,sha256=DfMdPleFFGADXLsr7zXui8BTz_y9futY6rNmNdh9y7k
 cache_dit/compile/utils.py,sha256=KU60xc474Anbj7Y_FLRFmNxEjVYLLXkhbtCLXO7o_Tc,3699
 cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cache_dit/custom_ops/triton_taylorseer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit-0.2.5.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.5.dist-info/METADATA,sha256=J37Waq-cMbuFfTrngXuxqouXpjHK9qhR_MZHlE2odmY,26249
-cache_dit-0.2.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.5.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.5.dist-info/RECORD,,
+cache_dit/metrics/__init__.py,sha256=5yavk_6b50EhaHij0C-7nQegOA0uD79-qX96dI1i_8s,461
+cache_dit/metrics/fid.py,sha256=RIC-wW0RFv6W9IW6nc6Ih4dAxunUTnzcugMbYZjdqRM,12891
+cache_dit/metrics/inception.py,sha256=pBVe2X6ylLPIXTG4-GWDM9DWnCviMJbJ45R3ulhktR0,12759
+cache_dit/metrics/metrics.py,sha256=FUrpc58ofg0LKyM3Y_7kfUVbevMfCFFcaeaxeOrj9iY,10498
+cache_dit-0.2.6.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.6.dist-info/METADATA,sha256=aTTyx_2gM7Z-mcIYd3_LlGljiZOwIq9iadDXOFPDzwA,27848
+cache_dit-0.2.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.6.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.6.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.6.dist-info/RECORD,,

cache_dit-0.2.6.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+cache-dit-metrics-cli = cache_dit.metrics:main