Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torchvision.datasets as datasets | |
| import torchvision.transforms as transforms | |
| import torchvision.transforms.functional as TF | |
| from torchvision.io import read_video | |
| from torch.utils.data import Dataset | |
| from random import random, choice, shuffle | |
| from io import BytesIO | |
| from PIL import Image | |
| from PIL import ImageFile | |
| from scipy.ndimage.filters import gaussian_filter | |
| import pickle | |
| import os | |
| MEAN = { | |
| "imagenet":[0.485, 0.456, 0.406], | |
| "clip":[0.48145466, 0.4578275, 0.40821073] | |
| } | |
| STD = { | |
| "imagenet":[0.229, 0.224, 0.225], | |
| "clip":[0.26862954, 0.26130258, 0.27577711] | |
| } | |
| def recursively_read(rootdir, must_contain, exts=["mp4", "avi"]): | |
| out = [] | |
| for r, d, f in os.walk(rootdir): | |
| for file in f: | |
| if (file.split('.')[1] in exts) and (must_contain in os.path.join(r, file)): | |
| out.append(os.path.join(r, file)) | |
| return out | |
| def get_list(path, must_contain=''): | |
| image_list = recursively_read(path, must_contain) | |
| return image_list | |
| def uniform_capture_video_frames(path, num_frames=16): | |
| capture = cv2.VideoCapture(path) | |
| total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| frame_interval = total_frames // num_frames | |
| frames = [] | |
| frame_count = 0 | |
| while len(frames) < num_frames: | |
| ret, frame = capture.read() | |
| if not ret: | |
| break | |
| if frame_count % frame_interval == 0: | |
| # 将OpenCV的BGR图像转换为RGB图像 | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| # 将numpy数组转换为PIL图像 | |
| pil_image = Image.fromarray(frame_rgb) | |
| frames.append(pil_image) | |
| frame_count += 1 | |
| capture.release() | |
| return frames | |
| class RealFakeDataset(Dataset): | |
| def __init__(self, opt, clip_model=None, transform=None, num_frames=16): | |
| self.opt = opt | |
| assert opt.data_label in ["train", "val"] | |
| self.data_label = opt.data_label | |
| self.num_frames = num_frames | |
| real_list = get_list( os.path.join(opt.real_list_path) ) | |
| fake_list = get_list( os.path.join(opt.fake_list_path) ) | |
| # setting the labels for the dataset | |
| self.labels_dict = {} | |
| for i in real_list: | |
| self.labels_dict[i] = 0 | |
| for i in fake_list: | |
| self.labels_dict[i] = 1 | |
| self.total_list = real_list + fake_list | |
| shuffle(self.total_list) | |
| self.transform = transform | |
| def __len__(self): | |
| return len(self.total_list) | |
| def __getitem__(self, idx): | |
| img_path = self.total_list[idx] | |
| label = self.labels_dict[img_path] | |
| # img = Image.open(img_path).convert("RGB") | |
| # video_frames = uniform_capture_video_frames(img_path, num_frames=32) | |
| # self.clip_model.to(torch.device('cuda:{}'.format(self.opt.gpu_ids[0])) if self.opt.gpu_ids else torch.device('cpu')) | |
| frames, _, _ = read_video(str(img_path), pts_unit='sec') | |
| frames = frames[:self.num_frames] | |
| frames = frames.permute(0, 3, 1, 2) # (T,H,W,C) -> (T,C,H,W) | |
| if self.transform is not None: | |
| video_frames = torch.cat([self.transform(TF.to_pil_image(frame)).unsqueeze(0) for frame in frames]) | |
| return video_frames, label |