From cf2e4f48d0abf5e0a98ad2151d982ebe5b4ee31b Mon Sep 17 00:00:00 2001 From: rain305f <2201212856@stu.pku.edu.cn> Date: Fri, 8 Mar 2024 09:11:05 +0000 Subject: [PATCH] [docs]:update the eval_code for calculate the FVD, clip_score, ssim, lpips, psnr Former-commit-id: 1ee226c338126fbdc5178020b12271b980a7ce0c --- eval_code/README.md | 233 +++++++++++++++++++ eval_code/ca_lpips.py | 97 ++++++++ eval_code/cal_fvd.py | 85 +++++++ eval_code/cal_psnr.py | 85 +++++++ eval_code/cal_ssim.py | 113 +++++++++ eval_code/clip_score.py | 226 ++++++++++++++++++ eval_code/fvd/styleganv/fvd.py | 90 +++++++ eval_code/fvd/videogpt/fvd.py | 137 +++++++++++ eval_code/fvd/videogpt/pytorch_i3d.py | 322 ++++++++++++++++++++++++++ 9 files changed, 1388 insertions(+) create mode 100644 eval_code/README.md create mode 100644 eval_code/ca_lpips.py create mode 100644 eval_code/cal_fvd.py create mode 100644 eval_code/cal_psnr.py create mode 100644 eval_code/cal_ssim.py create mode 100755 eval_code/clip_score.py create mode 100644 eval_code/fvd/styleganv/fvd.py create mode 100644 eval_code/fvd/videogpt/fvd.py create mode 100644 eval_code/fvd/videogpt/pytorch_i3d.py diff --git a/eval_code/README.md b/eval_code/README.md new file mode 100644 index 000000000..7b6f0a7df --- /dev/null +++ b/eval_code/README.md @@ -0,0 +1,233 @@ +# Metrics_on_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 +> ``` + +# Usage + + + +``` +# clip_score cross modality +python 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 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 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, 8 videos of a batch, 10 frames, 3 channels, 64x64 size. +import torch +from cal_fvd import calculate_fvd +from cal_psnr import calculate_psnr +from cal_ssim import calculate_ssim +from cal_lpips import calculate_lpips + +NUMBER_OF_VIDEOS = 8 +VIDEO_LENGTH = 30 +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 = {} +result['fvd'] = calculate_fvd(videos1, videos2, device, method='styleganv') +# result['fvd'] = calculate_fvd(videos1, videos2, device, method='videogpt') +result['ssim'] = calculate_ssim(videos1, videos2) +result['psnr'] = calculate_psnr(videos1, videos2) +result['lpips'] = calculate_lpips(videos1, videos2, device) +print(json.dumps(result, indent=4)) +``` + +It means we calculate: + +- `FVD-frames[:10]`, `FVD-frames[:11]`, ..., `FVD-frames[:30]` +- `avg-PSNR/SSIM/LPIPS-frame[0]`, `avg-PSNR/SSIM/LPIPS-frame[1]`, ..., `avg-PSNR/SSIM/LPIPS-frame[:30]`, and their std. + +We cannot calculate `FVD-frames[:8]`, and it will pass when calculating, see ps.6. + +The result shows: a all-zero matrix and a all-one matrix, their FVD-30 (FVD[:30]) is 151.17 (styleganv method). We also calculate their standard deviation. Other metrics are the same. And we use the calculation method of styleganv. + +``` +{ + "fvd": { + "value": { + "10": 570.07320378183, + "11": 486.1906542471159, + "12": 552.3373915075898, + "13": 146.6242330185728, + "14": 172.57268402948895, + "15": 133.88932632116126, + "16": 153.11023578170108, + "17": 357.56400892781204, + "18": 382.1335612721498, + "19": 306.7100176942531, + "20": 338.18221898178774, + "21": 77.95587603163293, + "22": 82.49997632357349, + "23": 64.41624523513073, + "24": 66.08097153313875, + "25": 314.4341061962642, + "26": 316.8616746151064, + "27": 288.884418528541, + "28": 287.8192683223724, + "29": 152.15076552354864, + "30": 151.16806952692093 + }, + "video_setting": [ + 8, + 3, + 30, + 64, + 64 + ], + "video_setting_name": "batch_size, channel, time, heigth, width" + }, + "video_setting": [ + 8, + 3, + 30, + 64, + 64 + ], + "video_setting_name": "batch_size, channel, time, heigth, width" + }, + "ssim": { + "value": { + "0": 9.999000099990664e-05, + ..., + "29": 9.999000099990664e-05 + }, + "value_std": { + "0": 0.0, + ..., + "29": 0.0 + }, + "video_setting": [ + 30, + 3, + 64, + 64 + ], + "video_setting_name": "time, channel, heigth, width" + }, + "psnr": { + "value": { + "0": 0.0, + ..., + "29": 0.0 + }, + "value_std": { + "0": 0.0, + ..., + "29": 0.0 + }, + "video_setting": [ + 30, + 3, + 64, + 64 + ], + "video_setting_name": "time, channel, heigth, width" + }, + "lpips": { + "value": { + "0": 0.8140146732330322, + ..., + "29": 0.8140146732330322 + }, + "value_std": { + "0": 0.0, + ..., + "29": 0.0 + }, + "video_setting": [ + 30, + 3, + 64, + 64 + ], + "video_setting_name": "time, channel, heigth, width" + } +} +``` + +# 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_code/ca_lpips.py b/eval_code/ca_lpips.py new file mode 100644 index 000000000..10edc610b --- /dev/null +++ b/eval_code/ca_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_code/cal_fvd.py b/eval_code/cal_fvd.py new file mode 100644 index 000000000..1f1a9806e --- /dev/null +++ b/eval_code/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_code/cal_psnr.py b/eval_code/cal_psnr.py new file mode 100644 index 000000000..b7a5ac023 --- /dev/null +++ b/eval_code/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_code/cal_ssim.py b/eval_code/cal_ssim.py new file mode 100644 index 000000000..b2de811b6 --- /dev/null +++ b/eval_code/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_code/clip_score.py b/eval_code/clip_score.py new file mode 100755 index 000000000..4972b814d --- /dev/null +++ b/eval_code/clip_score.py @@ -0,0 +1,226 @@ +"""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 + + +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')) + +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(): + 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_code/fvd/styleganv/fvd.py b/eval_code/fvd/styleganv/fvd.py new file mode 100644 index 000000000..3043a2a4a --- /dev/null +++ b/eval_code/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_code/fvd/videogpt/fvd.py b/eval_code/fvd/videogpt/fvd.py new file mode 100644 index 000000000..e81c82904 --- /dev/null +++ b/eval_code/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_code/fvd/videogpt/pytorch_i3d.py b/eval_code/fvd/videogpt/pytorch_i3d.py new file mode 100644 index 000000000..58a16cdd7 --- /dev/null +++ b/eval_code/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