opensora/eval/eval_common_metric.py

"""Calculates the CLIP Scores

The CLIP model is a contrasitively learned language-image model. There is
an image encoder and a text encoder. It is believed that the CLIP model could 
measure the similarity of cross modalities. Please find more information from 
https://github.com/openai/CLIP.

The CLIP Score measures the Cosine Similarity between two embedded features.
This repository utilizes the pretrained CLIP Model to calculate 
the mean average of cosine similarities. 

See --help to see further details.

Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.

Copyright 2023 The Hong Kong Polytechnic University

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import os.path as osp
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, Subset
from decord import VideoReader, cpu
import random
from pytorchvideo.transforms import ShortSideScale
from torchvision.io import read_video
from torchvision.transforms import Lambda, Compose
from torchvision.transforms._transforms_video import CenterCropVideo
import sys
sys.path.append(".")
from opensora.eval.cal_lpips import calculate_lpips
from opensora.eval.cal_fvd import calculate_fvd
from opensora.eval.cal_psnr import calculate_psnr
from opensora.eval.cal_flolpips import calculate_flolpips
from opensora.eval.cal_ssim import calculate_ssim

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

class VideoDataset(Dataset):
    def __init__(self, 
                 real_video_dir,
                 generated_video_dir,
                 num_frames,
                 sample_rate = 1,
                 crop_size=None,
                 resolution=128,
                 ) -> None:
        super().__init__()
        self.real_video_files = self._combine_without_prefix(real_video_dir)
        self.generated_video_files = self._combine_without_prefix(generated_video_dir)
        self.num_frames = num_frames
        self.sample_rate = sample_rate
        self.crop_size = crop_size
        self.short_size = resolution


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

    def __getitem__(self, index):
        if index >= len(self):
            raise IndexError
        real_video_file = self.real_video_files[index]
        generated_video_file = self.generated_video_files[index]
        print(real_video_file, generated_video_file)
        real_video_tensor  = self._load_video(real_video_file)
        generated_video_tensor  = self._load_video(generated_video_file)
        return {'real': real_video_tensor, 'generated':generated_video_tensor }


    def _load_video(self, video_path):
        num_frames = self.num_frames
        sample_rate = self.sample_rate
        decord_vr = VideoReader(video_path, ctx=cpu(0))
        total_frames = len(decord_vr)
        sample_frames_len = sample_rate * num_frames

        if total_frames >= sample_frames_len:
            s = 0
            e = s + sample_frames_len
            num_frames = num_frames
        else:
            s = 0
            e = total_frames
            num_frames = int(total_frames / sample_frames_len * num_frames)
            print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path,
                total_frames)


        frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
        video_data = decord_vr.get_batch(frame_id_list).asnumpy()
        video_data = torch.from_numpy(video_data)
        video_data = video_data.permute(0, 3, 1, 2) # (T, H, W, C) -> (C, T, H, W)
        return _preprocess(video_data, short_size=self.short_size, crop_size = self.crop_size)


    def _combine_without_prefix(self, folder_path, prefix='.'):
        folder = []
        os.makedirs(folder_path, exist_ok=True)
        for name in os.listdir(folder_path):
            if name[0] == prefix:
                continue
            if osp.isfile(osp.join(folder_path, name)):
                folder.append(osp.join(folder_path, name))
        folder.sort()
        return folder

def _preprocess(video_data, short_size=128, crop_size=None):
    transform = Compose(
        [
            Lambda(lambda x: x / 255.0),
            ShortSideScale(size=short_size),
            CenterCropVideo(crop_size=crop_size),
        ]
    )
    video_outputs = transform(video_data)
    # video_outputs = torch.unsqueeze(video_outputs, 0) # (bz,c,t,h,w)
    return video_outputs


def calculate_common_metric(args, dataloader, device):

    score_list = []
    for batch_data in tqdm(dataloader): # {'real': real_video_tensor, 'generated':generated_video_tensor }
        real_videos = batch_data['real'] 
        generated_videos = batch_data['generated']
        assert real_videos.shape[2] == generated_videos.shape[2]
        if args.metric == 'fvd':
            tmp_list = list(calculate_fvd(real_videos, generated_videos, args.device, method=args.fvd_method)['value'].values())
        elif args.metric == 'ssim':
            tmp_list = list(calculate_ssim(real_videos, generated_videos)['value'].values())
        elif args.metric == 'psnr':
            tmp_list = list(calculate_psnr(real_videos, generated_videos)['value'].values())
        elif args.metric == 'flolpips':
            result = calculate_flolpips(real_videos, generated_videos, args.device)
            tmp_list = list(result['value'].values())
        else:
            tmp_list  = list(calculate_lpips(real_videos, generated_videos, args.device)['value'].values())
        score_list += tmp_list
    return np.mean(score_list)
        
def main():
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('--batch_size', type=int, default=2,
                    help='Batch size to use')
    parser.add_argument('--real_video_dir', type=str,
                    help=('the path of real videos`'))
    parser.add_argument('--generated_video_dir', type=str,
                    help=('the path of generated videos`'))
    parser.add_argument('--device', type=str, default=None,
                    help='Device to use. Like cuda, cuda:0 or cpu')
    parser.add_argument('--num_workers', type=int, default=8,
                    help=('Number of processes to use for data loading. '
                          'Defaults to `min(8, num_cpus)`'))
    parser.add_argument('--sample_fps', type=int, default=30)
    parser.add_argument('--resolution', type=int, default=336)
    parser.add_argument('--crop_size', type=int, default=None)
    parser.add_argument('--num_frames', type=int, default=100)
    parser.add_argument('--sample_rate', type=int, default=1)
    parser.add_argument('--subset_size', type=int, default=None)
    parser.add_argument("--metric", type=str, default="fvd",choices=['fvd','psnr','ssim','lpips', 'flolpips'])
    parser.add_argument("--fvd_method", type=str, default='styleganv',choices=['styleganv','videogpt'])


    args = parser.parse_args()

    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:
            # os.sched_getaffinity is not available under Windows, use
            # os.cpu_count instead (which may not return the *available* number
            # of CPUs).
            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

    
    dataset = VideoDataset(args.real_video_dir,
                           args.generated_video_dir,
                            num_frames = args.num_frames,
                            sample_rate = args.sample_rate,
                            crop_size=args.crop_size,
                            resolution=args.resolution) 

    if args.subset_size:
        indices = range(args.subset_size)
        dataset = Subset(dataset, indices=indices)
    
    dataloader = DataLoader(dataset, args.batch_size, 
                            num_workers=num_workers, pin_memory=True)
    

    metric_score = calculate_common_metric(args, dataloader,device)
    print('metric: ', args.metric, " ",metric_score)

if __name__ == '__main__':
    main()