tools/metrics/clip-score/clip_score.py

import io
import os
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser

import clip
import numpy as np
import torch
import webdataset as wds
from PIL import Image
from torch.utils.data import DataLoader, Dataset, IterableDataset

from diffusion.data.transforms import get_transform
from tools.metrics.utils import tracker

try:
    from tqdm import tqdm
except ImportError:
    # If tqdm is not available, provide a mock version of it
    def tqdm(x):
        return x


import json

IMAGE_EXTENSIONS = {"bmp", "jpg", "jpeg", "pgm", "png", "ppm", "tif", "tiff", "webp"}
TEXT_EXTENSIONS = {"txt"}


class DummyDataset(Dataset):
    FLAGS = ["img", "txt", "json"]

    def __init__(
        self,
        real_path,
        fake_path,
        real_flag: str = "img",
        fake_flag: str = "img",
        gen_img_path="",
        transform=None,
        tokenizer=None,
    ) -> None:
        super().__init__()
        assert (
            real_flag in self.FLAGS and fake_flag in self.FLAGS
        ), f"CLIP Score only support modality of {self.FLAGS}. However, get {real_flag} and {fake_flag}"
        self.gen_img_path = gen_img_path
        print(f"images are from {gen_img_path}")
        self.real_folder = self._load_img_from_path(real_path)
        self.real_flag = real_flag
        self.fake_data = self._load_txt_from_path(fake_path)
        self.transform = transform
        self.tokenizer = tokenizer
        self.data_dict = {}

    def __len__(self):
        return len(self.real_folder)

    def __getitem__(self, index):
        if index >= len(self):
            raise IndexError
        real_path = self.real_folder[index]
        real_data = self._load_modality(real_path, self.real_flag)
        fake_data = self._load_txt(self.fake_data[index])
        sample = dict(real=real_data, fake=fake_data, prompt=self.fake_data[index])
        return sample

    def _load_modality(self, path, modality):
        if modality == "img":
            data = self._load_img(path)
        else:
            raise TypeError(f"Got unexpected modality: {modality}")
        return data

    def _load_txt(self, data):
        if self.tokenizer is not None:
            data = self.tokenizer(data, context_length=77, truncate=True).squeeze()
        return data

    def _load_img(self, path):
        img = Image.open(path)
        if self.transform is not None:
            img = self.transform(img)
        return img

    def _load_img_from_path(self, path):
        image_list = []
        if path.endswith(".json"):
            with open(path) as file:
                data_dict = json.load(file)
            all_lines = list(data_dict.keys())[:sample_nums]
            if isinstance(all_lines, list):
                for k in all_lines:
                    img_path = os.path.join(self.gen_img_path, f"{k}.jpg")
                    image_list.append(img_path)
            elif isinstance(all_lines, dict):
                assert sample_nums >= 30_000, ValueError(f"{sample_nums} is not supported for json files")
                for k, v in all_lines.items():
                    img_path = os.path.join(self.gen_img_path, f"{k}.jpg")
                    image_list.append(img_path)

        else:
            raise ValueError(f"Only JSON file type is supported now. Wrong with: {path}")

        return image_list

    def _load_txt_from_path(self, path):
        txt_list = []
        if path.endswith(".json"):
            with open(path) as file:
                data_dict = json.load(file)
            all_lines = list(data_dict.keys())[:sample_nums]
            if isinstance(all_lines, list):
                for k in all_lines:
                    v = data_dict[k]
                    txt_list.append(v["prompt"])
            elif isinstance(all_lines, dict):
                assert sample_nums >= 30_000, ValueError(f"{sample_nums} is not supported for json files")
                for k, v in all_lines.items():
                    txt_list.append(v["prompt"])
        else:
            raise ValueError(f"Only JSON file type is supported now. Wrong with: {path}")

        return txt_list


class DummyTarDataset(IterableDataset):
    def __init__(
        self, tar_path, transform=None, external_json_path=None, prompt_key="prompt", tokenizer=None, **kwargs
    ):
        assert ".tar" in tar_path
        self.sample_nums = args.sample_nums
        self.dataset = (
            wds.WebDataset(tar_path)
            .map(self.safe_decode)
            .to_tuple("png;jpg", "json", "__key__")
            .map(self.process_sample)
            .slice(0, self.sample_nums)
        )
        if external_json_path is not None and os.path.exists(external_json_path):
            print(f"Loading {external_json_path}, wait...")
            self.json_file = json.load(open(external_json_path))
        else:
            self.json_file = {}
            assert prompt_key == "prompt"
        self.prompt_key = prompt_key
        self.transform = transform
        self.tokenizer = tokenizer

    def __iter__(self):
        return self._generator()

    def _generator(self):
        for i, (ori_img, info, key) in enumerate(self.dataset):
            if self.transform is not None:
                img = self.transform(ori_img)

            if key in self.json_file:
                info.update(self.json_file[key])

            prompt = info.get(self.prompt_key, "")
            if not prompt:
                prompt = ""
                print(f"{self.prompt_key} not exist in {key}.json")
            txt_feat = self._load_txt(prompt)

            yield dict(
                real=img, fake=txt_feat, prompt=prompt, ori_img=np.array(img), key=key, prompt_key=self.prompt_key
            )

    def __len__(self):
        return self.sample_nums

    def _load_txt(self, data):
        if self.tokenizer is not None:
            data = self.tokenizer(data, context_length=77, truncate=True).squeeze()
        return data

    @staticmethod
    def process_sample(sample):
        try:
            image_bytes, json_bytes, key = sample
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            json_dict = json.loads(json_bytes)
            return image, json_dict, key
        except (ValueError, TypeError, OSError) as e:
            print(f"Skipping sample due to error: {e}")
            return None

    @staticmethod
    def safe_decode(sample):
        def custom_decode(sample):
            result = {}
            for k, v in sample.items():
                result[k] = v
            return result

        try:
            return custom_decode(sample)
        except Exception as e:
            print(f"skipping sample due to decode error: {e}")
            return None


@torch.no_grad()
def calculate_clip_score(dataloader, model, real_flag, fake_flag, save_json_path=None):
    score_acc = 0.0
    sample_num = 0.0
    json_dict = {} if save_json_path is not None else None
    logit_scale = model.logit_scale.exp()
    for batch_data in tqdm(dataloader, desc=f"CLIP-Score: {args.exp_name}", position=args.gpu_id, leave=True):
        real_features = forward_modality(model, batch_data["real"], real_flag)
        fake_features = forward_modality(model, batch_data["fake"], fake_flag)

        # normalize features
        real_features = real_features / real_features.norm(dim=1, keepdim=True).to(torch.float32)
        fake_features = fake_features / fake_features.norm(dim=1, keepdim=True).to(torch.float32)

        score = logit_scale * (fake_features * real_features).sum()
        if save_json_path is not None:
            json_dict[batch_data["key"][0]] = {f"{batch_data['prompt_key'][0]}": f"{score:.04f}"}

        score_acc += score
        sample_num += batch_data["real"].shape[0]

    if save_json_path is not None:
        json.dump(json_dict, open(save_json_path, "w"))
    return score_acc / sample_num


@torch.no_grad()
def calculate_clip_score_official(dataloader):
    import numpy as np
    from torchmetrics.multimodal.clip_score import CLIPScore

    clip_score_fn = CLIPScore(model_name_or_path="openai/clip-vit-large-patch14").to(device)
    # clip_score_fn = CLIPScore(model_name_or_path="openai/clip-vit-base-patch16").to(device)
    all_clip_scores = []

    for batch_data in tqdm(dataloader, desc=args.exp_name, position=args.gpu_id, leave=True):
        imgs = batch_data["real"].add_(1.0).mul_(0.5)
        imgs = (imgs * 255).to(dtype=torch.uint8, device=device)

        prompts = batch_data["prompt"]
        clip_scores = clip_score_fn(imgs, prompts).detach().cpu()
        all_clip_scores.append(float(clip_scores))

    clip_scores = float(np.mean(all_clip_scores))
    return clip_scores


def forward_modality(model, data, flag):
    device = next(model.parameters()).device
    if flag == "img":
        features = model.encode_image(data.to(device))
    elif flag == "txt":
        features = model.encode_text(data.to(device))
    else:
        raise TypeError
    return features


def main():
    txt_path = args.txt_path if args.txt_path is not None else args.img_path
    gen_img_path = str(os.path.join(args.img_path, args.exp_name))
    if ".tar" in gen_img_path:
        save_txt_path = os.path.join(txt_path, f"{args.exp_name}_{args.tar_prompt_key}_clip_score.txt").replace(
            ".tar", ""
        )
        save_json_path = save_txt_path.replace(".tar", "").replace(".txt", ".json")
        if os.path.exists(save_json_path):
            print(f"{save_json_path} exists. Finished.")
            return None
    else:
        save_txt_path = os.path.join(txt_path, f"{args.exp_name}_sample{sample_nums}_clip_score.txt")
        save_json_path = None
    if os.path.exists(save_txt_path):
        with open(save_txt_path) as f:
            clip_score = f.readlines()[0].strip()
        print(f"CLIP Score:  {clip_score}: {args.exp_name}")
        return {args.exp_name: float(clip_score)}

    print(f"Loading CLIP model: {args.clip_model}")
    if args.clipscore_type == "diffusers":
        preprocess = get_transform("default_train", 512)
    else:
        model, preprocess = clip.load(args.clip_model, device=device)

    if ".tar" in gen_img_path:
        dataset = DummyTarDataset(
            gen_img_path,
            transform=preprocess,
            external_json_path=args.external_json_file,
            prompt_key=args.tar_prompt_key,
            tokenizer=clip.tokenize,
        )
    else:
        dataset = DummyDataset(
            args.real_path,
            args.fake_path,
            args.real_flag,
            args.fake_flag,
            transform=preprocess,
            tokenizer=clip.tokenize,
            gen_img_path=gen_img_path,
        )
    dataloader = DataLoader(dataset, args.batch_size, num_workers=num_workers, pin_memory=True)

    print("Calculating CLIP Score:")
    if args.clipscore_type == "diffusers":
        clip_score = calculate_clip_score_official(dataloader)
    else:
        clip_score = calculate_clip_score(
            dataloader, model, args.real_flag, args.fake_flag, save_json_path=save_json_path
        )
        clip_score = clip_score.cpu().item()
    print("CLIP Score: ", clip_score)
    with open(save_txt_path, "w") as file:
        file.write(str(clip_score))
    print(f"Result saved at: {save_txt_path}")

    return {args.exp_name: clip_score}


def parse_args():
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument("--batch-size", type=int, default=50, help="Batch size to use")
    parser.add_argument("--clip-model", type=str, default="ViT-L/14", help="CLIP model to use")
    # parser.add_argument('--clip-model', type=str, default='ViT-B/16', help='CLIP model to use')
    parser.add_argument("--img_path", type=str, default=None)
    parser.add_argument("--txt_path", type=str, default=None)
    parser.add_argument("--sample_nums", type=int, default=30_000)
    parser.add_argument("--exp_name", type=str, default="Sana")
    parser.add_argument(
        "--num-workers", type=int, help="Number of processes to use for data loading.  Defaults to `min(8, num_cpus)`"
    )
    parser.add_argument("--device", type=str, default=None, help="Device to use. Like cuda, cuda:0 or cpu")
    parser.add_argument("--real_flag", type=str, default="img", help="The modality of real path. Default to img")
    parser.add_argument("--fake_flag", type=str, default="txt", help="The modality of real path. Default to txt")
    parser.add_argument("--real_path", type=str, help="Paths to the generated images")
    parser.add_argument("--fake_path", type=str, help="Paths to the generated images")
    parser.add_argument("--external_json_file", type=str, default=None, help="external meta json file for tar_file")
    parser.add_argument("--tar_prompt_key", type=str, default="prompt", help="key name of prompt in json")

    # online logging setting
    parser.add_argument("--clipscore_type", type=str, default="self", choices=["diffusers", "self"])
    parser.add_argument("--log_metric", type=str, default="metric")
    parser.add_argument("--gpu_id", type=int, default=0)
    parser.add_argument("--log_clip_score", action="store_true")
    parser.add_argument("--suffix_label", type=str, default="", help="used for clip_score online log")
    parser.add_argument("--tracker_pattern", type=str, default="epoch_step", help="used for fid online log")
    parser.add_argument(
        "--report_to",
        type=str,
        default=None,
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
        ),
    )
    parser.add_argument(
        "--tracker_project_name",
        type=str,
        default="t2i-evit-baseline",
        help=(
            "The `project_name` argument passed to Accelerator.init_trackers for"
            " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"
        ),
    )
    parser.add_argument(
        "--name",
        type=str,
        default="baseline",
        help=("Wandb Project Name"),
    )
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()
    sample_nums = args.sample_nums
    if args.device is None:
        device = torch.device("cuda" if (torch.cuda.is_available()) else "cpu")
    else:
        device = torch.device(args.device)

    if args.num_workers is None:
        try:
            num_cpus = len(os.sched_getaffinity(0))
        except AttributeError:
            num_cpus = os.cpu_count()
        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
    else:
        num_workers = args.num_workers

    args.exp_name = os.path.basename(args.exp_name) or os.path.dirname(args.exp_name)
    clip_score_result = main()
    if args.log_clip_score:
        tracker(args, clip_score_result, args.suffix_label, pattern=args.tracker_pattern, metric="CLIP-Score")