diff --git a/docs/EVAL_README.md b/docs/EVAL_README.md new file mode 100644 index 000000000..df8be8d36 --- /dev/null +++ b/docs/EVAL_README.md @@ -0,0 +1,152 @@ +# Evaluate the Generated Video_Quality + +You can easily calculate the following video quality metrics, which supports the batch-wise process. +- **CLIP-SCORE**: It uses the pretrained CLIP model to measure the cosine similarity between two modalities. +- **FVD**: Frechét Video Distance +- **SSIM**: structural similarity index measure +- **LPIPS**: learned perceptual image patch similarity +- **PSNR**: peak-signal-to-noise ratio + +# Requrirement +## Environemnt +- install Pytorch (torch>=1.7.1) +- install CLIP + ``` + pip install git+https://github.com/openai/CLIP.git + ``` +- install clip-cose from PyPi + ``` + pip install clip-score + ``` +- Other package + ``` + pip install lpips + pip install scipy (scipy==1.7.3/1.9.3, if you use 1.11.3, **you will calculate a WRONG FVD VALUE!!!**) + pip install numpy + pip install pillow + pip install torchvision>=0.8.2 + pip install ftfy + pip install regex + pip install tqdm + ``` +## Pretrain model +- FVD + Before you cacluate FVD, you should first download the FVD pre-trained model. You can manually download any of the following and put it into FVD folder. + - `i3d_torchscript.pt` from [here](https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt) + - `i3d_pretrained_400.pt` from [here](https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI) + + +## Other Notices +1. Make sure the pixel value of videos should be in [0, 1]. +2. We average SSIM when images have 3 channels, ssim is the only metric extremely sensitive to gray being compared to b/w. +3. Because the i3d model downsamples in the time dimension, `frames_num` should > 10 when calculating FVD, so FVD calculation begins from 10-th frame, like upper example. +4. For grayscale videos, we multiply to 3 channels +5. data input specifications for clip_score +> - Image Files:All images should be stored in a single directory. The image files can be in either .png or .jpg format. +> +> - Text Files: All text data should be contained in plain text files in a separate directory. These text files should have the extension .txt. +> +> Note: The number of files in the image directory should be exactly equal to the number of files in the text directory. Additionally, the files in the image directory and text directory should be paired by file name. For instance, if there is a cat.png in the image directory, there should be a corresponding cat.txt in the text directory. +> +> Directory Structure Example: +> ``` +> ├── path/to/image +> │ ├── cat.png +> │ ├── dog.png +> │ └── bird.jpg +> └── path/to/text +> ├── cat.txt +> ├── dog.txt +> └── bird.txt +> ``` + +6. data input specifications for fvd, psnr, ssim, lpips + +> Directory Structure Example: +> ``` +> ├── path/to/generated_image +> │ ├── cat.mp4 +> │ ├── dog.mp4 +> │ └── bird.mp4 +> └── path/to/real_image +> ├── cat.mp4 +> ├── dog.mp4 +> └── bird.mp4 +> ``` + + + +# Usage + + + +``` +# clip_score cross modality +python eval_clip_score.py \ + --real_path path/to/image \ + --generated_path path/to/text \ + --batch-size 50 \ + --device "cuda" + +# clip_score within the same modality +python eval_clip_score.py \ + --real_path path/to/textA \ + --generated_path path/to/textB \ + --real_flag txt \ + --generated_flag txt \ + --batch-size 50 \ + --device "cuda" + +python eval_clip_score.py \ + --real_path path/to/imageA \ + --generated_path path/to/imageB \ + --real_flag img \ + --generated_flag img \ + --batch-size 50 \ + --device "cuda" + + +# fvd, psnr, ssim, lpips, you need to set the frame_num, resolution etc... + +python eval_common_metric.py \ + --real_video_dir path/to/imageA\ + --generated_video_dir path/to/imageB \ + --batch_size 10 \ + --crop_size 64 \ + --num_frames 20 \ + --device 'cuda' \ + --metric 'fvd' \ + --fvd_method 'styleganv' + +python eval_common_metric.py \ + --real_video_dir path/to/imageA\ + --generated_video_dir path/to/imageB \ + --batch_size 10 \ + --num_frames 20 \ + --crop_size 64 \ + --device 'cuda' \ + --metric 'psnr' + + +python eval_common_metric.py \ + --real_video_dir path/to/imageA\ + --generated_video_dir path/to/imageB \ + --batch_size 10 \ + --num_frames 20 \ + --crop_size 64 \ + --device 'cuda' \ + --metric 'ssim' + +python eval_common_metric.py \ + --real_video_dir path/to/imageA\ + --generated_video_dir path/to/imageB \ + --batch_size 10 \ + --num_frames 20 \ + --crop_size 64 \ + --device 'cuda' \ + --metric 'lpips' + +``` + +# Acknowledgement +The evaluation codebase refers to [clip-score](https://github.com/Taited/clip-score) and [common_metrics](https://github.com/JunyaoHu/common_metrics_on_video_quality). \ No newline at end of file diff --git a/eval/cal_clip_score.sh b/eval/cal_clip_score.sh new file mode 100644 index 000000000..19a2c462d --- /dev/null +++ b/eval/cal_clip_score.sh @@ -0,0 +1,37 @@ +python eval_common_metric.py \ + --real_video_dir /data/xiaogeng_liu/data/video1 \ + --generated_video_dir /data/xiaogeng_liu/data/video2 \ + --batch_size 10 \ + --crop_size 64 \ + --num_frames 20 \ + --device 'cuda' \ + --metric 'fvd' \ + --fvd_method 'styleganv' + +python eval_common_metric.py \ + --real_video_dir /data/xiaogeng_liu/data/video1 \ + --generated_video_dir /data/xiaogeng_liu/data/video2 \ + --batch_size 10 \ + --num_frames 20 \ + --crop_size 64 \ + --device 'cuda' \ + --metric 'psnr' + + +python eval_common_metric.py \ + --real_video_dir /data/xiaogeng_liu/data/video1 \ + --generated_video_dir /data/xiaogeng_liu/data/video2 \ + --batch_size 10 \ + --num_frames 20 \ + --crop_size 64 \ + --device 'cuda' \ + --metric 'ssim' + +python eval_common_metric.py \ + --real_video_dir /data/xiaogeng_liu/data/video1 \ + --generated_video_dir /data/xiaogeng_liu/data/video2 \ + --batch_size 10 \ + --num_frames 20 \ + --crop_size 64 \ + --device 'cuda' \ + --metric 'lpips' \ No newline at end of file diff --git a/eval/cal_fvd.py b/eval/cal_fvd.py new file mode 100644 index 000000000..1f1a9806e --- /dev/null +++ b/eval/cal_fvd.py @@ -0,0 +1,85 @@ +import numpy as np +import torch +from tqdm import tqdm + +def trans(x): + # if greyscale images add channel + if x.shape[-3] == 1: + x = x.repeat(1, 1, 3, 1, 1) + + # permute BTCHW -> BCTHW + x = x.permute(0, 2, 1, 3, 4) + + return x + +def calculate_fvd(videos1, videos2, device, method='styleganv'): + + if method == 'styleganv': + from fvd.styleganv.fvd import get_fvd_feats, frechet_distance, load_i3d_pretrained + elif method == 'videogpt': + from fvd.videogpt.fvd import load_i3d_pretrained + from fvd.videogpt.fvd import get_fvd_logits as get_fvd_feats + from fvd.videogpt.fvd import frechet_distance + + print("calculate_fvd...") + + # videos [batch_size, timestamps, channel, h, w] + + assert videos1.shape == videos2.shape + + i3d = load_i3d_pretrained(device=device) + fvd_results = [] + + # support grayscale input, if grayscale -> channel*3 + # BTCHW -> BCTHW + # videos -> [batch_size, channel, timestamps, h, w] + + videos1 = trans(videos1) + videos2 = trans(videos2) + + fvd_results = {} + + # for calculate FVD, each clip_timestamp must >= 10 + for clip_timestamp in tqdm(range(10, videos1.shape[-3]+1)): + + # get a video clip + # videos_clip [batch_size, channel, timestamps[:clip], h, w] + videos_clip1 = videos1[:, :, : clip_timestamp] + videos_clip2 = videos2[:, :, : clip_timestamp] + + # get FVD features + feats1 = get_fvd_feats(videos_clip1, i3d=i3d, device=device) + feats2 = get_fvd_feats(videos_clip2, i3d=i3d, device=device) + + # calculate FVD when timestamps[:clip] + fvd_results[clip_timestamp] = frechet_distance(feats1, feats2) + + result = { + "value": fvd_results, + "video_setting": videos1.shape, + "video_setting_name": "batch_size, channel, time, heigth, width", + } + + return result + +# test code / using example + +def main(): + NUMBER_OF_VIDEOS = 8 + VIDEO_LENGTH = 50 + CHANNEL = 3 + SIZE = 64 + videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + device = torch.device("cuda") + # device = torch.device("cpu") + + import json + result = calculate_fvd(videos1, videos2, device, method='videogpt') + print(json.dumps(result, indent=4)) + + result = calculate_fvd(videos1, videos2, device, method='styleganv') + print(json.dumps(result, indent=4)) + +if __name__ == "__main__": + main() diff --git a/eval/cal_lpips.py b/eval/cal_lpips.py new file mode 100644 index 000000000..10edc610b --- /dev/null +++ b/eval/cal_lpips.py @@ -0,0 +1,97 @@ +import numpy as np +import torch +from tqdm import tqdm +import math + +import torch +import lpips + +spatial = True # Return a spatial map of perceptual distance. + +# Linearly calibrated models (LPIPS) +loss_fn = lpips.LPIPS(net='alex', spatial=spatial) # Can also set net = 'squeeze' or 'vgg' +# loss_fn = lpips.LPIPS(net='alex', spatial=spatial, lpips=False) # Can also set net = 'squeeze' or 'vgg' + +def trans(x): + # if greyscale images add channel + if x.shape[-3] == 1: + x = x.repeat(1, 1, 3, 1, 1) + + # value range [0, 1] -> [-1, 1] + x = x * 2 - 1 + + return x + +def calculate_lpips(videos1, videos2, device): + # image should be RGB, IMPORTANT: normalized to [-1,1] + print("calculate_lpips...") + + assert videos1.shape == videos2.shape + + # videos [batch_size, timestamps, channel, h, w] + + # support grayscale input, if grayscale -> channel*3 + # value range [0, 1] -> [-1, 1] + videos1 = trans(videos1) + videos2 = trans(videos2) + + lpips_results = [] + + for video_num in tqdm(range(videos1.shape[0])): + # get a video + # video [timestamps, channel, h, w] + video1 = videos1[video_num] + video2 = videos2[video_num] + + lpips_results_of_a_video = [] + for clip_timestamp in range(len(video1)): + # get a img + # img [timestamps[x], channel, h, w] + # img [channel, h, w] tensor + + img1 = video1[clip_timestamp].unsqueeze(0).to(device) + img2 = video2[clip_timestamp].unsqueeze(0).to(device) + + loss_fn.to(device) + + # calculate lpips of a video + lpips_results_of_a_video.append(loss_fn.forward(img1, img2).mean().detach().cpu().tolist()) + lpips_results.append(lpips_results_of_a_video) + + lpips_results = np.array(lpips_results) + + lpips = {} + lpips_std = {} + + for clip_timestamp in range(len(video1)): + lpips[clip_timestamp] = np.mean(lpips_results[:,clip_timestamp]) + lpips_std[clip_timestamp] = np.std(lpips_results[:,clip_timestamp]) + + + result = { + "value": lpips, + "value_std": lpips_std, + "video_setting": video1.shape, + "video_setting_name": "time, channel, heigth, width", + } + + return result + +# test code / using example + +def main(): + NUMBER_OF_VIDEOS = 8 + VIDEO_LENGTH = 50 + CHANNEL = 3 + SIZE = 64 + videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + device = torch.device("cuda") + # device = torch.device("cpu") + + import json + result = calculate_lpips(videos1, videos2, device) + print(json.dumps(result, indent=4)) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/eval/cal_psnr.py b/eval/cal_psnr.py new file mode 100644 index 000000000..b7a5ac023 --- /dev/null +++ b/eval/cal_psnr.py @@ -0,0 +1,85 @@ +import numpy as np +import torch +from tqdm import tqdm +import math + +def img_psnr(img1, img2): + # [0,1] + # compute mse + # mse = np.mean((img1-img2)**2) + mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2) + # compute psnr + if mse < 1e-10: + return 100 + psnr = 20 * math.log10(1 / math.sqrt(mse)) + return psnr + +def trans(x): + return x + +def calculate_psnr(videos1, videos2): + print("calculate_psnr...") + + # videos [batch_size, timestamps, channel, h, w] + + assert videos1.shape == videos2.shape + + videos1 = trans(videos1) + videos2 = trans(videos2) + + psnr_results = [] + + for video_num in tqdm(range(videos1.shape[0])): + # get a video + # video [timestamps, channel, h, w] + video1 = videos1[video_num] + video2 = videos2[video_num] + + psnr_results_of_a_video = [] + for clip_timestamp in range(len(video1)): + # get a img + # img [timestamps[x], channel, h, w] + # img [channel, h, w] numpy + + img1 = video1[clip_timestamp].numpy() + img2 = video2[clip_timestamp].numpy() + + # calculate psnr of a video + psnr_results_of_a_video.append(img_psnr(img1, img2)) + + psnr_results.append(psnr_results_of_a_video) + + psnr_results = np.array(psnr_results) + + psnr = {} + psnr_std = {} + + for clip_timestamp in range(len(video1)): + psnr[clip_timestamp] = np.mean(psnr_results[:,clip_timestamp]) + psnr_std[clip_timestamp] = np.std(psnr_results[:,clip_timestamp]) + + result = { + "value": psnr, + "value_std": psnr_std, + "video_setting": video1.shape, + "video_setting_name": "time, channel, heigth, width", + } + + return result + +# test code / using example + +def main(): + NUMBER_OF_VIDEOS = 8 + VIDEO_LENGTH = 50 + CHANNEL = 3 + SIZE = 64 + videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + + import json + result = calculate_psnr(videos1, videos2) + print(json.dumps(result, indent=4)) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/eval/cal_ssim.py b/eval/cal_ssim.py new file mode 100644 index 000000000..b2de811b6 --- /dev/null +++ b/eval/cal_ssim.py @@ -0,0 +1,113 @@ +import numpy as np +import torch +from tqdm import tqdm +import cv2 + +def ssim(img1, img2): + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + img1 = img1.astype(np.float64) + img2 = img2.astype(np.float64) + kernel = cv2.getGaussianKernel(11, 1.5) + window = np.outer(kernel, kernel.transpose()) + mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid + mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5] + mu1_sq = mu1 ** 2 + mu2_sq = mu2 ** 2 + mu1_mu2 = mu1 * mu2 + sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq + sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq + sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2 + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * + (sigma1_sq + sigma2_sq + C2)) + return ssim_map.mean() + + +def calculate_ssim_function(img1, img2): + # [0,1] + # ssim is the only metric extremely sensitive to gray being compared to b/w + if not img1.shape == img2.shape: + raise ValueError('Input images must have the same dimensions.') + if img1.ndim == 2: + return ssim(img1, img2) + elif img1.ndim == 3: + if img1.shape[0] == 3: + ssims = [] + for i in range(3): + ssims.append(ssim(img1[i], img2[i])) + return np.array(ssims).mean() + elif img1.shape[0] == 1: + return ssim(np.squeeze(img1), np.squeeze(img2)) + else: + raise ValueError('Wrong input image dimensions.') + +def trans(x): + return x + +def calculate_ssim(videos1, videos2): + print("calculate_ssim...") + + # videos [batch_size, timestamps, channel, h, w] + + assert videos1.shape == videos2.shape + + videos1 = trans(videos1) + videos2 = trans(videos2) + + ssim_results = [] + + for video_num in tqdm(range(videos1.shape[0])): + # get a video + # video [timestamps, channel, h, w] + video1 = videos1[video_num] + video2 = videos2[video_num] + + ssim_results_of_a_video = [] + for clip_timestamp in range(len(video1)): + # get a img + # img [timestamps[x], channel, h, w] + # img [channel, h, w] numpy + + img1 = video1[clip_timestamp].numpy() + img2 = video2[clip_timestamp].numpy() + + # calculate ssim of a video + ssim_results_of_a_video.append(calculate_ssim_function(img1, img2)) + + ssim_results.append(ssim_results_of_a_video) + + ssim_results = np.array(ssim_results) + + ssim = {} + ssim_std = {} + + for clip_timestamp in range(len(video1)): + ssim[clip_timestamp] = np.mean(ssim_results[:,clip_timestamp]) + ssim_std[clip_timestamp] = np.std(ssim_results[:,clip_timestamp]) + + result = { + "value": ssim, + "value_std": ssim_std, + "video_setting": video1.shape, + "video_setting_name": "time, channel, heigth, width", + } + + return result + +# test code / using example + +def main(): + NUMBER_OF_VIDEOS = 8 + VIDEO_LENGTH = 50 + CHANNEL = 3 + SIZE = 64 + videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False) + device = torch.device("cuda") + + import json + result = calculate_ssim(videos1, videos2) + print(json.dumps(result, indent=4)) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/eval/eval_clip_score.py b/eval/eval_clip_score.py new file mode 100755 index 000000000..ed2f51082 --- /dev/null +++ b/eval/eval_clip_score.py @@ -0,0 +1,225 @@ +"""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 clip +import torch +from PIL import Image +from torch.utils.data import Dataset, DataLoader + +try: + from tqdm import tqdm +except ImportError: + # If tqdm is not available, provide a mock version of it + def tqdm(x): + return x + + +IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm', + 'tif', 'tiff', 'webp'} + +TEXT_EXTENSIONS = {'txt'} + + +class DummyDataset(Dataset): + + FLAGS = ['img', 'txt'] + def __init__(self, real_path, generated_path, + real_flag: str = 'img', + generated_flag: str = 'img', + transform = None, + tokenizer = None) -> None: + super().__init__() + assert real_flag in self.FLAGS and generated_flag in self.FLAGS, \ + 'CLIP Score only support modality of {}. However, get {} and {}'.format( + self.FLAGS, real_flag, generated_flag + ) + self.real_folder = self._combine_without_prefix(real_path) + self.real_flag = real_flag + self.fake_foler = self._combine_without_prefix(generated_path) + self.generated_flag = generated_flag + self.transform = transform + self.tokenizer = tokenizer + # assert self._check() + + def __len__(self): + return len(self.real_folder) + + def __getitem__(self, index): + if index >= len(self): + raise IndexError + real_path = self.real_folder[index] + generated_path = self.fake_foler[index] + real_data = self._load_modality(real_path, self.real_flag) + fake_data = self._load_modality(generated_path, self.generated_flag) + + sample = dict(real=real_data, fake=fake_data) + return sample + + def _load_modality(self, path, modality): + if modality == 'img': + data = self._load_img(path) + elif modality == 'txt': + data = self._load_txt(path) + else: + raise TypeError("Got unexpected modality: {}".format(modality)) + 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_txt(self, path): + with open(path, 'r') as fp: + data = fp.read() + fp.close() + if self.tokenizer is not None: + data = self.tokenizer(data).squeeze() + return data + + def _check(self): + for idx in range(len(self)): + real_name = self.real_folder[idx].split('.') + fake_name = self.fake_folder[idx].split('.') + if fake_name != real_name: + return False + return True + + def _combine_without_prefix(self, folder_path, prefix='.'): + folder = [] + for name in os.listdir(folder_path): + if name[0] == prefix: + continue + folder.append(osp.join(folder_path, name)) + folder.sort() + return folder + + +@torch.no_grad() +def calculate_clip_score(dataloader, model, real_flag, generated_flag): + score_acc = 0. + sample_num = 0. + logit_scale = model.logit_scale.exp() + for batch_data in tqdm(dataloader): + real = batch_data['real'] + real_features = forward_modality(model, real, real_flag) + fake = batch_data['fake'] + fake_features = forward_modality(model, fake, generated_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) + + # calculate scores + # score = logit_scale * real_features @ fake_features.t() + # score_acc += torch.diag(score).sum() + score = logit_scale * (fake_features * real_features).sum() + score_acc += score + sample_num += real.shape[0] + + return score_acc / sample_num + + +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(): + 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-B/32', + help='CLIP model to use') + 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('--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('--generated_flag', type=str, default='txt', + help=('The modality of generated path. ' + 'Default to txt')) + parser.add_argument('--real_path', type=str, + help=('Paths to the real images or ' + 'to .npz statistic files')) + parser.add_argument('--generated_path', type=str, + help=('Paths to the generated images or ' + 'to .npz statistic files')) + 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 + + print('Loading CLIP model: {}'.format(args.clip_model)) + model, preprocess = clip.load(args.clip_model, device=device) + + dataset = DummyDataset(args.real_path, args.generated_path, + args.real_flag, args.generated_flag, + transform=preprocess, tokenizer=clip.tokenize) + dataloader = DataLoader(dataset, args.batch_size, + num_workers=num_workers, pin_memory=True) + + print('Calculating CLIP Score:') + clip_score = calculate_clip_score(dataloader, model, + args.real_flag, args.generated_flag) + clip_score = clip_score.cpu().item() + print('CLIP Score: ', clip_score) + + +if __name__ == '__main__': + main() diff --git a/eval/eval_common_metric.py b/eval/eval_common_metric.py new file mode 100644 index 000000000..67df44aa4 --- /dev/null +++ b/eval/eval_common_metric.py @@ -0,0 +1,213 @@ +"""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 clip +import torch +from PIL import Image +from torch.utils.data import Dataset, DataLoader +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 RandomCropVideo +from eval.cal_lpips import calculate_lpips +from eval.cal_fvd import calculate_fvd +from eval.cal_psnr import calculate_psnr +from 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] + 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 = random.randint(0, total_frames - sample_frames_len - 1) + 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 = [] + for name in os.listdir(folder_path): + if name[0] == prefix: + continue + 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) - 0.5)), + ShortSideScale(size=short_size), + RandomCropVideo(size=crop_size) if crop_size is not None else Lambda(lambda x: x), + + ] + ) + 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'] + + if args.metric == 'fvd': + tmp_list = list(calculate_fvd(real_videos, generated_videos, args.device, method=args.fvd_method)['value'].keys()) + elif args.metric == 'ssim': + tmp_list = list(calculate_ssim(real_videos, generated_videos)['value'].keys()) + elif args.metric == 'psnr': + tmp_list = list(calculate_psnr(real_videos, generated_videos)['value'].keys()) + else: + tmp_list = list(calculate_lpips(real_videos, generated_videos, args.device)['value'].keys()) + 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("--metric", type=str, default="fvd",choices=['fvd','psnr','ssim','lpips']) + 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) + + 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() diff --git a/eval/fvd/styleganv/fvd.py b/eval/fvd/styleganv/fvd.py new file mode 100644 index 000000000..3043a2a4a --- /dev/null +++ b/eval/fvd/styleganv/fvd.py @@ -0,0 +1,90 @@ +import torch +import os +import math +import torch.nn.functional as F + +# https://github.com/universome/fvd-comparison + + +def load_i3d_pretrained(device=torch.device('cpu')): + i3D_WEIGHTS_URL = "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt" + filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_torchscript.pt') + print(filepath) + if not os.path.exists(filepath): + print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.") + os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}") + i3d = torch.jit.load(filepath).eval().to(device) + i3d = torch.nn.DataParallel(i3d) + return i3d + + +def get_feats(videos, detector, device, bs=10): + # videos : torch.tensor BCTHW [0, 1] + detector_kwargs = dict(rescale=False, resize=False, return_features=True) # Return raw features before the softmax layer. + feats = np.empty((0, 400)) + with torch.no_grad(): + for i in range((len(videos)-1)//bs + 1): + feats = np.vstack([feats, detector(torch.stack([preprocess_single(video) for video in videos[i*bs:(i+1)*bs]]).to(device), **detector_kwargs).detach().cpu().numpy()]) + return feats + + +def get_fvd_feats(videos, i3d, device, bs=10): + # videos in [0, 1] as torch tensor BCTHW + # videos = [preprocess_single(video) for video in videos] + embeddings = get_feats(videos, i3d, device, bs) + return embeddings + + +def preprocess_single(video, resolution=224, sequence_length=None): + # video: CTHW, [0, 1] + c, t, h, w = video.shape + + # temporal crop + if sequence_length is not None: + assert sequence_length <= t + video = video[:, :sequence_length] + + # scale shorter side to resolution + scale = resolution / min(h, w) + if h < w: + target_size = (resolution, math.ceil(w * scale)) + else: + target_size = (math.ceil(h * scale), resolution) + video = F.interpolate(video, size=target_size, mode='bilinear', align_corners=False) + + # center crop + c, t, h, w = video.shape + w_start = (w - resolution) // 2 + h_start = (h - resolution) // 2 + video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution] + + # [0, 1] -> [-1, 1] + video = (video - 0.5) * 2 + + return video.contiguous() + + +""" +Copy-pasted from https://github.com/cvpr2022-stylegan-v/stylegan-v/blob/main/src/metrics/frechet_video_distance.py +""" +from typing import Tuple +from scipy.linalg import sqrtm +import numpy as np + + +def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + mu = feats.mean(axis=0) # [d] + sigma = np.cov(feats, rowvar=False) # [d, d] + return mu, sigma + + +def frechet_distance(feats_fake: np.ndarray, feats_real: np.ndarray) -> float: + mu_gen, sigma_gen = compute_stats(feats_fake) + mu_real, sigma_real = compute_stats(feats_real) + m = np.square(mu_gen - mu_real).sum() + if feats_fake.shape[0]>1: + s, _ = sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member + fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2)) + else: + fid = np.real(m) + return float(fid) \ No newline at end of file diff --git a/eval/fvd/videogpt/fvd.py b/eval/fvd/videogpt/fvd.py new file mode 100644 index 000000000..e81c82904 --- /dev/null +++ b/eval/fvd/videogpt/fvd.py @@ -0,0 +1,137 @@ +import torch +import os +import math +import torch.nn.functional as F +import numpy as np +import einops + +def load_i3d_pretrained(device=torch.device('cpu')): + i3D_WEIGHTS_URL = "https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI" + filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_pretrained_400.pt') + print(filepath) + if not os.path.exists(filepath): + print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.") + os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}") + from .pytorch_i3d import InceptionI3d + i3d = InceptionI3d(400, in_channels=3).eval().to(device) + i3d.load_state_dict(torch.load(filepath, map_location=device)) + i3d = torch.nn.DataParallel(i3d) + return i3d + +def preprocess_single(video, resolution, sequence_length=None): + # video: THWC, {0, ..., 255} + video = video.permute(0, 3, 1, 2).float() / 255. # TCHW + t, c, h, w = video.shape + + # temporal crop + if sequence_length is not None: + assert sequence_length <= t + video = video[:sequence_length] + + # scale shorter side to resolution + scale = resolution / min(h, w) + if h < w: + target_size = (resolution, math.ceil(w * scale)) + else: + target_size = (math.ceil(h * scale), resolution) + video = F.interpolate(video, size=target_size, mode='bilinear', + align_corners=False) + + # center crop + t, c, h, w = video.shape + w_start = (w - resolution) // 2 + h_start = (h - resolution) // 2 + video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution] + video = video.permute(1, 0, 2, 3).contiguous() # CTHW + + video -= 0.5 + + return video + +def preprocess(videos, target_resolution=224): + # we should tras videos in [0-1] [b c t h w] as th.float + # -> videos in {0, ..., 255} [b t h w c] as np.uint8 array + videos = einops.rearrange(videos, 'b c t h w -> b t h w c') + videos = (videos*255).numpy().astype(np.uint8) + + b, t, h, w, c = videos.shape + videos = torch.from_numpy(videos) + videos = torch.stack([preprocess_single(video, target_resolution) for video in videos]) + return videos * 2 # [-0.5, 0.5] -> [-1, 1] + +def get_fvd_logits(videos, i3d, device, bs=10): + videos = preprocess(videos) + embeddings = get_logits(i3d, videos, device, bs=10) + return embeddings + +# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L161 +def _symmetric_matrix_square_root(mat, eps=1e-10): + u, s, v = torch.svd(mat) + si = torch.where(s < eps, s, torch.sqrt(s)) + return torch.matmul(torch.matmul(u, torch.diag(si)), v.t()) + +# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L400 +def trace_sqrt_product(sigma, sigma_v): + sqrt_sigma = _symmetric_matrix_square_root(sigma) + sqrt_a_sigmav_a = torch.matmul(sqrt_sigma, torch.matmul(sigma_v, sqrt_sigma)) + return torch.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a)) + +# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2 +def cov(m, rowvar=False): + '''Estimate a covariance matrix given data. + + Covariance indicates the level to which two variables vary together. + If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`, + then the covariance matrix element `C_{ij}` is the covariance of + `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`. + + Args: + m: A 1-D or 2-D array containing multiple variables and observations. + Each row of `m` represents a variable, and each column a single + observation of all those variables. + rowvar: If `rowvar` is True, then each row represents a + variable, with observations in the columns. Otherwise, the + relationship is transposed: each column represents a variable, + while the rows contain observations. + + Returns: + The covariance matrix of the variables. + ''' + if m.dim() > 2: + raise ValueError('m has more than 2 dimensions') + if m.dim() < 2: + m = m.view(1, -1) + if not rowvar and m.size(0) != 1: + m = m.t() + + fact = 1.0 / (m.size(1) - 1) # unbiased estimate + m -= torch.mean(m, dim=1, keepdim=True) + mt = m.t() # if complex: mt = m.t().conj() + return fact * m.matmul(mt).squeeze() + + +def frechet_distance(x1, x2): + x1 = x1.flatten(start_dim=1) + x2 = x2.flatten(start_dim=1) + m, m_w = x1.mean(dim=0), x2.mean(dim=0) + sigma, sigma_w = cov(x1, rowvar=False), cov(x2, rowvar=False) + mean = torch.sum((m - m_w) ** 2) + if x1.shape[0]>1: + sqrt_trace_component = trace_sqrt_product(sigma, sigma_w) + trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component + fd = trace + mean + else: + fd = np.real(mean) + return float(fd) + + +def get_logits(i3d, videos, device, bs=10): + # assert videos.shape[0] % 16 == 0 + with torch.no_grad(): + logits = [] + for i in range(0, videos.shape[0], bs): + batch = videos[i:i + bs].to(device) + # logits.append(i3d.module.extract_features(batch)) # wrong + logits.append(i3d(batch)) # right + logits = torch.cat(logits, dim=0) + return logits diff --git a/eval/fvd/videogpt/pytorch_i3d.py b/eval/fvd/videogpt/pytorch_i3d.py new file mode 100644 index 000000000..58a16cdd7 --- /dev/null +++ b/eval/fvd/videogpt/pytorch_i3d.py @@ -0,0 +1,322 @@ +# Original code from https://github.com/piergiaj/pytorch-i3d +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np + +class MaxPool3dSamePadding(nn.MaxPool3d): + + def compute_pad(self, dim, s): + if s % self.stride[dim] == 0: + return max(self.kernel_size[dim] - self.stride[dim], 0) + else: + return max(self.kernel_size[dim] - (s % self.stride[dim]), 0) + + def forward(self, x): + # compute 'same' padding + (batch, channel, t, h, w) = x.size() + out_t = np.ceil(float(t) / float(self.stride[0])) + out_h = np.ceil(float(h) / float(self.stride[1])) + out_w = np.ceil(float(w) / float(self.stride[2])) + pad_t = self.compute_pad(0, t) + pad_h = self.compute_pad(1, h) + pad_w = self.compute_pad(2, w) + + pad_t_f = pad_t // 2 + pad_t_b = pad_t - pad_t_f + pad_h_f = pad_h // 2 + pad_h_b = pad_h - pad_h_f + pad_w_f = pad_w // 2 + pad_w_b = pad_w - pad_w_f + + pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b) + x = F.pad(x, pad) + return super(MaxPool3dSamePadding, self).forward(x) + + +class Unit3D(nn.Module): + + def __init__(self, in_channels, + output_channels, + kernel_shape=(1, 1, 1), + stride=(1, 1, 1), + padding=0, + activation_fn=F.relu, + use_batch_norm=True, + use_bias=False, + name='unit_3d'): + + """Initializes Unit3D module.""" + super(Unit3D, self).__init__() + + self._output_channels = output_channels + self._kernel_shape = kernel_shape + self._stride = stride + self._use_batch_norm = use_batch_norm + self._activation_fn = activation_fn + self._use_bias = use_bias + self.name = name + self.padding = padding + + self.conv3d = nn.Conv3d(in_channels=in_channels, + out_channels=self._output_channels, + kernel_size=self._kernel_shape, + stride=self._stride, + padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function + bias=self._use_bias) + + if self._use_batch_norm: + self.bn = nn.BatchNorm3d(self._output_channels, eps=1e-5, momentum=0.001) + + def compute_pad(self, dim, s): + if s % self._stride[dim] == 0: + return max(self._kernel_shape[dim] - self._stride[dim], 0) + else: + return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0) + + + def forward(self, x): + # compute 'same' padding + (batch, channel, t, h, w) = x.size() + out_t = np.ceil(float(t) / float(self._stride[0])) + out_h = np.ceil(float(h) / float(self._stride[1])) + out_w = np.ceil(float(w) / float(self._stride[2])) + pad_t = self.compute_pad(0, t) + pad_h = self.compute_pad(1, h) + pad_w = self.compute_pad(2, w) + + pad_t_f = pad_t // 2 + pad_t_b = pad_t - pad_t_f + pad_h_f = pad_h // 2 + pad_h_b = pad_h - pad_h_f + pad_w_f = pad_w // 2 + pad_w_b = pad_w - pad_w_f + + pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b) + x = F.pad(x, pad) + + x = self.conv3d(x) + if self._use_batch_norm: + x = self.bn(x) + if self._activation_fn is not None: + x = self._activation_fn(x) + return x + + + +class InceptionModule(nn.Module): + def __init__(self, in_channels, out_channels, name): + super(InceptionModule, self).__init__() + + self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0, + name=name+'/Branch_0/Conv3d_0a_1x1') + self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0, + name=name+'/Branch_1/Conv3d_0a_1x1') + self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3], + name=name+'/Branch_1/Conv3d_0b_3x3') + self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0, + name=name+'/Branch_2/Conv3d_0a_1x1') + self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3], + name=name+'/Branch_2/Conv3d_0b_3x3') + self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3], + stride=(1, 1, 1), padding=0) + self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0, + name=name+'/Branch_3/Conv3d_0b_1x1') + self.name = name + + def forward(self, x): + b0 = self.b0(x) + b1 = self.b1b(self.b1a(x)) + b2 = self.b2b(self.b2a(x)) + b3 = self.b3b(self.b3a(x)) + return torch.cat([b0,b1,b2,b3], dim=1) + + +class InceptionI3d(nn.Module): + """Inception-v1 I3D architecture. + The model is introduced in: + Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset + Joao Carreira, Andrew Zisserman + https://arxiv.org/pdf/1705.07750v1.pdf. + See also the Inception architecture, introduced in: + Going deeper with convolutions + Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, + Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich. + http://arxiv.org/pdf/1409.4842v1.pdf. + """ + + # Endpoints of the model in order. During construction, all the endpoints up + # to a designated `final_endpoint` are returned in a dictionary as the + # second return value. + VALID_ENDPOINTS = ( + 'Conv3d_1a_7x7', + 'MaxPool3d_2a_3x3', + 'Conv3d_2b_1x1', + 'Conv3d_2c_3x3', + 'MaxPool3d_3a_3x3', + 'Mixed_3b', + 'Mixed_3c', + 'MaxPool3d_4a_3x3', + 'Mixed_4b', + 'Mixed_4c', + 'Mixed_4d', + 'Mixed_4e', + 'Mixed_4f', + 'MaxPool3d_5a_2x2', + 'Mixed_5b', + 'Mixed_5c', + 'Logits', + 'Predictions', + ) + + def __init__(self, num_classes=400, spatial_squeeze=True, + final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5): + """Initializes I3D model instance. + Args: + num_classes: The number of outputs in the logit layer (default 400, which + matches the Kinetics dataset). + spatial_squeeze: Whether to squeeze the spatial dimensions for the logits + before returning (default True). + final_endpoint: The model contains many possible endpoints. + `final_endpoint` specifies the last endpoint for the model to be built + up to. In addition to the output at `final_endpoint`, all the outputs + at endpoints up to `final_endpoint` will also be returned, in a + dictionary. `final_endpoint` must be one of + InceptionI3d.VALID_ENDPOINTS (default 'Logits'). + name: A string (optional). The name of this module. + Raises: + ValueError: if `final_endpoint` is not recognized. + """ + + if final_endpoint not in self.VALID_ENDPOINTS: + raise ValueError('Unknown final endpoint %s' % final_endpoint) + + super(InceptionI3d, self).__init__() + self._num_classes = num_classes + self._spatial_squeeze = spatial_squeeze + self._final_endpoint = final_endpoint + self.logits = None + + if self._final_endpoint not in self.VALID_ENDPOINTS: + raise ValueError('Unknown final endpoint %s' % self._final_endpoint) + + self.end_points = {} + end_point = 'Conv3d_1a_7x7' + self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7], + stride=(2, 2, 2), padding=(3,3,3), name=name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'MaxPool3d_2a_3x3' + self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2), + padding=0) + if self._final_endpoint == end_point: return + + end_point = 'Conv3d_2b_1x1' + self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0, + name=name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Conv3d_2c_3x3' + self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1, + name=name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'MaxPool3d_3a_3x3' + self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2), + padding=0) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_3b' + self.end_points[end_point] = InceptionModule(192, [64,96,128,16,32,32], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_3c' + self.end_points[end_point] = InceptionModule(256, [128,128,192,32,96,64], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'MaxPool3d_4a_3x3' + self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2), + padding=0) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_4b' + self.end_points[end_point] = InceptionModule(128+192+96+64, [192,96,208,16,48,64], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_4c' + self.end_points[end_point] = InceptionModule(192+208+48+64, [160,112,224,24,64,64], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_4d' + self.end_points[end_point] = InceptionModule(160+224+64+64, [128,128,256,24,64,64], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_4e' + self.end_points[end_point] = InceptionModule(128+256+64+64, [112,144,288,32,64,64], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_4f' + self.end_points[end_point] = InceptionModule(112+288+64+64, [256,160,320,32,128,128], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'MaxPool3d_5a_2x2' + self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2), + padding=0) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_5b' + self.end_points[end_point] = InceptionModule(256+320+128+128, [256,160,320,32,128,128], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Mixed_5c' + self.end_points[end_point] = InceptionModule(256+320+128+128, [384,192,384,48,128,128], name+end_point) + if self._final_endpoint == end_point: return + + end_point = 'Logits' + self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7], + stride=(1, 1, 1)) + self.dropout = nn.Dropout(dropout_keep_prob) + self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes, + kernel_shape=[1, 1, 1], + padding=0, + activation_fn=None, + use_batch_norm=False, + use_bias=True, + name='logits') + + self.build() + + + def replace_logits(self, num_classes): + self._num_classes = num_classes + self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes, + kernel_shape=[1, 1, 1], + padding=0, + activation_fn=None, + use_batch_norm=False, + use_bias=True, + name='logits') + + + def build(self): + for k in self.end_points.keys(): + self.add_module(k, self.end_points[k]) + + def forward(self, x): + for end_point in self.VALID_ENDPOINTS: + if end_point in self.end_points: + x = self._modules[end_point](x) # use _modules to work with dataparallel + + x = self.logits(self.dropout(self.avg_pool(x))) + if self._spatial_squeeze: + logits = x.squeeze(3).squeeze(3) + logits = logits.mean(dim=2) + # logits is batch X time X classes, which is what we want to work with + return logits + + + def extract_features(self, x): + for end_point in self.VALID_ENDPOINTS: + if end_point in self.end_points: + x = self._modules[end_point](x) + return self.avg_pool(x) \ No newline at end of file