import os import json import numpy as np import logging import subprocess import torch from PIL import Image, ImageSequence from decord import VideoReader, cpu from torchvision import transforms from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, ToPILImage try: from torchvision.transforms import InterpolationMode BICUBIC = InterpolationMode.BICUBIC BILINEAR = InterpolationMode.BILINEAR except ImportError: BICUBIC = Image.BICUBIC BILINEAR = Image.BILINEAR CACHE_DIR = os.environ.get('VBENCH_CACHE_DIR') if CACHE_DIR is None: CACHE_DIR = os.path.join(os.path.expanduser('~'), '.cache', 'vbench') logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def clip_transform(n_px): return Compose([ Resize(n_px, interpolation=BICUBIC), CenterCrop(n_px), transforms.Lambda(lambda x: x.float().div(255.0)), Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), ]) def clip_transform_Image(n_px): return Compose([ Resize(n_px, interpolation=BICUBIC), CenterCrop(n_px), ToTensor(), Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), ]) def dino_transform(n_px): return Compose([ Resize(size=n_px), transforms.Lambda(lambda x: x.float().div(255.0)), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) def dino_transform_Image(n_px): return Compose([ Resize(size=n_px), ToTensor(), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) def tag2text_transform(n_px): normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) return Compose([ToPILImage(),Resize((n_px, n_px)),ToTensor(),normalize]) def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1): if sample in ["rand", "middle"]: # uniform sampling acc_samples = min(num_frames, vlen) # split the video into `acc_samples` intervals, and sample from each interval. intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int) ranges = [] for idx, interv in enumerate(intervals[:-1]): ranges.append((interv, intervals[idx + 1] - 1)) if sample == 'rand': try: frame_indices = [random.choice(range(x[0], x[1])) for x in ranges] except: frame_indices = np.random.permutation(vlen)[:acc_samples] frame_indices.sort() frame_indices = list(frame_indices) elif fix_start is not None: frame_indices = [x[0] + fix_start for x in ranges] elif sample == 'middle': frame_indices = [(x[0] + x[1]) // 2 for x in ranges] else: raise NotImplementedError if len(frame_indices) < num_frames: # padded with last frame padded_frame_indices = [frame_indices[-1]] * num_frames padded_frame_indices[:len(frame_indices)] = frame_indices frame_indices = padded_frame_indices elif "fps" in sample: # fps0.5, sequentially sample frames at 0.5 fps output_fps = float(sample[3:]) duration = float(vlen) / input_fps delta = 1 / output_fps # gap between frames, this is also the clip length each frame represents frame_seconds = np.arange(0 + delta / 2, duration + delta / 2, delta) frame_indices = np.around(frame_seconds * input_fps).astype(int) frame_indices = [e for e in frame_indices if e < vlen] if max_num_frames > 0 and len(frame_indices) > max_num_frames: frame_indices = frame_indices[:max_num_frames] # frame_indices = np.linspace(0 + delta / 2, duration + delta / 2, endpoint=False, num=max_num_frames) else: raise ValueError return frame_indices def load_video(video_path, data_transform=None, num_frames=None, return_tensor=True, width=None, height=None): """ Load a video from a given path and apply optional data transformations. The function supports loading video in GIF (.gif), PNG (.png), and MP4 (.mp4) formats. Depending on the format, it processes and extracts frames accordingly. Parameters: - video_path (str): The file path to the video or image to be loaded. - data_transform (callable, optional): A function that applies transformations to the video data. Returns: - frames (torch.Tensor): A tensor containing the video frames with shape (T, C, H, W), where T is the number of frames, C is the number of channels, H is the height, and W is the width. Raises: - NotImplementedError: If the video format is not supported. The function first determines the format of the video file by its extension. For GIFs, it iterates over each frame and converts them to RGB. For PNGs, it reads the single frame, converts it to RGB. For MP4s, it reads the frames using the VideoReader class and converts them to NumPy arrays. If a data_transform is provided, it is applied to the buffer before converting it to a tensor. Finally, the tensor is permuted to match the expected (T, C, H, W) format. """ if video_path.endswith('.gif'): frame_ls = [] img = Image.open(video_path) for frame in ImageSequence.Iterator(img): frame = frame.convert('RGB') frame = np.array(frame).astype(np.uint8) frame_ls.append(frame) buffer = np.array(frame_ls).astype(np.uint8) elif video_path.endswith('.png'): frame = Image.open(video_path) frame = frame.convert('RGB') frame = np.array(frame).astype(np.uint8) frame_ls = [frame] buffer = np.array(frame_ls) elif video_path.endswith('.mp4'): import decord decord.bridge.set_bridge('native') if width: video_reader = VideoReader(video_path, width=width, height=height, num_threads=1) else: video_reader = VideoReader(video_path, num_threads=1) frames = video_reader.get_batch(range(len(video_reader))) # (T, H, W, C), torch.uint8 buffer = frames.asnumpy().astype(np.uint8) else: raise NotImplementedError frames = buffer if num_frames: frame_indices = get_frame_indices( num_frames, len(frames), sample="middle" ) frames = frames[frame_indices] if data_transform: frames = data_transform(frames) elif return_tensor: frames = torch.Tensor(frames) frames = frames.permute(0, 3, 1, 2) # (T, C, H, W), torch.uint8 return frames def read_frames_decord_by_fps( video_path, sample_fps=2, sample='rand', fix_start=None, max_num_frames=-1, trimmed30=False, num_frames=8 ): import decord decord.bridge.set_bridge("torch") video_reader = VideoReader(video_path, num_threads=1) vlen = len(video_reader) fps = video_reader.get_avg_fps() duration = vlen / float(fps) if trimmed30 and duration > 30: duration = 30 vlen = int(30 * float(fps)) frame_indices = get_frame_indices( num_frames, vlen, sample=sample, fix_start=fix_start, input_fps=fps, max_num_frames=max_num_frames ) frames = video_reader.get_batch(frame_indices) # (T, H, W, C), torch.uint8 frames = frames.permute(0, 3, 1, 2) # (T, C, H, W), torch.uint8 return frames def load_dimension_info(json_dir, dimension, lang): """ Load video list and prompt information based on a specified dimension and language from a JSON file. Parameters: - json_dir (str): The directory path where the JSON file is located. - dimension (str): The dimension for evaluation to filter the video prompts. - lang (str): The language key used to retrieve the appropriate prompt text. Returns: - video_list (list): A list of video file paths that match the specified dimension. - prompt_dict_ls (list): A list of dictionaries, each containing a prompt and its corresponding video list. The function reads the JSON file to extract video information. It filters the prompts based on the specified dimension and compiles a list of video paths and associated prompts in the specified language. Notes: - The JSON file is expected to contain a list of dictionaries with keys 'dimension', 'video_list', and language-based prompts. - The function assumes that the 'video_list' key in the JSON can either be a list or a single string value. """ video_list = [] prompt_dict_ls = [] full_prompt_list = load_json(json_dir) for prompt_dict in full_prompt_list: if dimension in prompt_dict['dimension'] and 'video_list' in prompt_dict: prompt = prompt_dict[f'prompt_{lang}'] cur_video_list = prompt_dict['video_list'] if isinstance(prompt_dict['video_list'], list) else [prompt_dict['video_list']] video_list += cur_video_list if 'auxiliary_info' in prompt_dict and dimension in prompt_dict['auxiliary_info']: prompt_dict_ls += [{'prompt': prompt, 'video_list': cur_video_list, 'auxiliary_info': prompt_dict['auxiliary_info'][dimension]}] else: prompt_dict_ls += [{'prompt': prompt, 'video_list': cur_video_list}] return video_list, prompt_dict_ls def init_submodules(dimension_list, local=False, read_frame=False): submodules_dict = {} if local: logger.info("\x1b[32m[Local Mode]\x1b[0m Working in local mode, please make sure that the pre-trained model has been fully downloaded.") for dimension in dimension_list: os.makedirs(CACHE_DIR, exist_ok=True) if dimension == 'background_consistency': # read_frame = False if local: vit_b_path = f'{CACHE_DIR}/clip_model/ViT-B-32.pt' if not os.path.isfile(vit_b_path): wget_command = ['wget', 'https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt', '-P', os.path.dirname(vit_b_path)] subprocess.run(wget_command, check=True) else: vit_b_path = 'ViT-B/32' submodules_dict[dimension] = [vit_b_path, read_frame] elif dimension == 'human_action': umt_path = f'{CACHE_DIR}/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth' if not os.path.isfile(umt_path): wget_command = ['wget', 'https://pjlab-gvm-data.oss-cn-shanghai.aliyuncs.com/umt/single_modality/l16_ptk710_ftk710_ftk400_f16_res224.pth', '-P', os.path.dirname(umt_path)] subprocess.run(wget_command, check=True) submodules_dict[dimension] = [umt_path,] elif dimension == 'temporal_flickering': submodules_dict[dimension] = [] elif dimension == 'motion_smoothness': CUR_DIR = os.path.dirname(os.path.abspath(__file__)) submodules_dict[dimension] = { 'config': f'{CUR_DIR}/third_party/amt/cfgs/AMT-S.yaml', 'ckpt': f'{CACHE_DIR}/amt_model/amt-s.pth' } details = submodules_dict[dimension] # Check if the file exists, if not, download it with wget if not os.path.isfile(details['ckpt']): print(f"File {details['ckpt']} does not exist. Downloading...") wget_command = ['wget', '-P', os.path.dirname(details['ckpt']), 'https://huggingface.co/lalala125/AMT/resolve/main/amt-s.pth'] subprocess.run(wget_command, check=True) elif dimension == 'dynamic_degree': submodules_dict[dimension] = { 'model': f'{CACHE_DIR}/raft_model/models/raft-things.pth' } details = submodules_dict[dimension] if not os.path.isfile(details['model']): # raise NotImplementedError print(f"File {details['model']} does not exist. Downloading...") wget_command = ['wget', '-P', f'{CACHE_DIR}/raft_model/', 'https://dl.dropboxusercontent.com/s/4j4z58wuv8o0mfz/models.zip'] unzip_command = ['unzip', '-d', f'{CACHE_DIR}/raft_model/', f'{CACHE_DIR}/raft_model/models.zip'] remove_command = ['rm', '-r', f'{CACHE_DIR}/raft_model/models.zip'] try: subprocess.run(wget_command, check=True) subprocess.run(unzip_command, check=True) subprocess.run(remove_command, check=True) except subprocess.CalledProcessError as err: print(f"Error during downloading RAFT model: {err}") # Assign the DINO model path for subject consistency dimension elif dimension == 'subject_consistency': if local: submodules_dict[dimension] = { 'repo_or_dir': f'{CACHE_DIR}/dino_model/facebookresearch_dino_main/', 'path': f'{CACHE_DIR}/dino_model/dino_vitbase16_pretrain.pth', 'model': 'dino_vitb16', 'source': 'local', 'read_frame': read_frame } details = submodules_dict[dimension] # Check if the file exists, if not, download it with wget if not os.path.isdir(details['repo_or_dir']): print(f"Directory {details['repo_or_dir']} does not exist. Cloning repository...") subprocess.run(['git', 'clone', 'https://github.com/facebookresearch/dino', details['repo_or_dir']], check=True) if not os.path.isfile(details['path']): print(f"File {details['path']} does not exist. Downloading...") wget_command = ['wget', '-P', os.path.dirname(details['path']), 'https://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth'] subprocess.run(wget_command, check=True) else: submodules_dict[dimension] = { 'repo_or_dir':'facebookresearch/dino:main', 'source':'github', 'model': 'dino_vitb16', 'read_frame': read_frame } elif dimension == 'aesthetic_quality': aes_path = f'{CACHE_DIR}/aesthetic_model/emb_reader' if local: vit_l_path = f'{CACHE_DIR}/clip_model/ViT-L-14.pt' if not os.path.isfile(vit_l_path): wget_command = ['wget' ,'https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt', '-P', os.path.dirname(vit_l_path)] subprocess.run(wget_command, check=True) else: vit_l_path = 'ViT-L/14' submodules_dict[dimension] = [vit_l_path, aes_path] elif dimension == 'imaging_quality': musiq_spaq_path = f'{CACHE_DIR}/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth' if not os.path.isfile(musiq_spaq_path): wget_command = ['wget', 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/musiq_spaq_ckpt-358bb6af.pth', '-P', os.path.dirname(musiq_spaq_path)] subprocess.run(wget_command, check=True) submodules_dict[dimension] = {'model_path': musiq_spaq_path} elif dimension in ["object_class", "multiple_objects", "color", "spatial_relationship" ]: submodules_dict[dimension] = { "model_weight": f'{CACHE_DIR}/grit_model/grit_b_densecap_objectdet.pth' } if not os.path.exists(submodules_dict[dimension]['model_weight']): wget_command = ['wget', 'https://datarelease.blob.core.windows.net/grit/models/grit_b_densecap_objectdet.pth', '-P', os.path.dirname(submodules_dict[dimension]["model_weight"])] subprocess.run(wget_command, check=True) elif dimension == 'scene': submodules_dict[dimension] = { "pretrained": f'{CACHE_DIR}/caption_model/tag2text_swin_14m.pth', "image_size":384, "vit":"swin_b" } if not os.path.exists(submodules_dict[dimension]['pretrained']): wget_command = ['wget', 'https://huggingface.co/spaces/xinyu1205/recognize-anything/resolve/main/tag2text_swin_14m.pth', '-P', os.path.dirname(submodules_dict[dimension]["pretrained"])] subprocess.run(wget_command, check=True) elif dimension in ['appearance_style', 'culture_fairness']: if local: submodules_dict[dimension] = {"name": f'{CACHE_DIR}/clip_model/ViT-B-32.pt'} if not os.path.isfile(submodules_dict[dimension]["name"]): wget_command = ['wget', 'https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt', '-P', os.path.dirname(submodules_dict[dimension]["name"])] subprocess.run(wget_command, check=True) else: submodules_dict[dimension] = {"name": 'ViT-B/32'} elif dimension in ["temporal_style", "overall_consistency"]: submodules_dict[dimension] = { "pretrain": f'{CACHE_DIR}/ViCLIP/ViClip-InternVid-10M-FLT.pth', } if not os.path.exists(submodules_dict[dimension]['pretrain']): wget_command = ['wget', 'https://pjlab-gvm-data.oss-cn-shanghai.aliyuncs.com/internvideo/viclip/ViClip-InternVid-10M-FLT.pth', '-P', os.path.dirname(submodules_dict[dimension]["pretrain"])] subprocess.run(wget_command, check=True) elif dimension in ["gender_bias", "skin_bias"]: if local: submodules_dict[dimension] = { "name": f'{CACHE_DIR}/clip_model/ViT-B-32.pt', "retina": f'{CACHE_DIR}/retina_face_model/retinaface_resnet50_2020-07-20-f168fae3c.zip' } if not os.path.isfile(submodules_dict[dimension]["name"]): wget_command = ['wget', 'https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt', '-P', os.path.dirname(submodules_dict[dimension]["name"])] subprocess.run(wget_command, check=True) if not os.path.isfile(submodules_dict[dimension]["retina"]): wget_command = ['wget', 'https://github.com/ternaus/retinaface/releases/download/0.01/retinaface_resnet50_2020-07-20-f168fae3c.zip', '-P', os.path.dirname(submodules_dict[dimension]["retina"])] subprocess.run(wget_command, check=True) else: submodules_dict[dimension] = { "name": 'ViT-B/32', "retina": 'https://github.com/ternaus/retinaface/releases/download/0.01/retinaface_resnet50_2020-07-20-f168fae3c.zip'} elif dimension == 'safety': if local: submodules_dict[dimension] = { "name": f'{CACHE_DIR}/clip_model/ViT-B-32.pt', "sd_checker": f'{CACHE_DIR}/sd_safety_checker/', "q16": f'{CACHE_DIR}/q16/prompts.p' } if not os.path.isfile(submodules_dict[dimension]["name"]): wget_command = ['wget', 'https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt', '-P', os.path.dirname(submodules_dict[dimension]["name"])] subprocess.run(wget_command, check=True) if not os.path.isdir(submodules_dict[dimension]["sd_checker"]): wget_command_1 = ['wget', 'https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/config.json', '-P', submodules_dict[dimension]["sd_checker"]] wget_command_2 = ['wget', 'https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/preprocessor_config.json', '-P', submodules_dict[dimension]["sd_checker"]] wget_command_3 = ['wget', 'https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/pytorch_model.bin', '-P', submodules_dict[dimension]["sd_checker"]] subprocess.run(wget_command_1, check=True) subprocess.run(wget_command_2, check=True) subprocess.run(wget_command_3, check=True) if not os.path.isfile(submodules_dict[dimension]["q16"]): wget_command = ['wget', 'https://raw.githubusercontent.com/ml-research/Q16/main/data/ViT-B-32/prompts.p', '-P', os.path.dirname(submodules_dict[dimension]["q16"])] subprocess.run(wget_command, check=True) else: submodules_dict[dimension] = { "name": 'ViT-B/32', "sd_checker": 'CompVis/stable-diffusion-safety-checker', "q16": 'https://raw.githubusercontent.com/ml-research/Q16/main/data/ViT-B-32/prompts.p'} return submodules_dict def save_json(data, path, indent=4): with open(path, 'w', encoding='utf-8') as f: json.dump(data, f, indent=indent) def load_json(path): """ Load a JSON file from the given file path. Parameters: - file_path (str): The path to the JSON file. Returns: - data (dict or list): The data loaded from the JSON file, which could be a dictionary or a list. """ with open(path, 'r', encoding='utf-8') as f: return json.load(f)