Spaces:
Running
Running
File size: 4,220 Bytes
94bafa8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
from .sky_datasets import Sky
from torchvision import transforms
from .taichi_datasets import Taichi
from datasets import video_transforms
from .ucf101_datasets import UCF101
from .ffs_datasets import FaceForensics
from .ffs_image_datasets import FaceForensicsImages
from .sky_image_datasets import SkyImages
from .ucf101_image_datasets import UCF101Images
from .taichi_image_datasets import TaichiImages
def get_dataset(args):
temporal_sample = video_transforms.TemporalRandomCrop(args.num_frames * args.frame_interval) # 16 1
if args.dataset == 'ffs':
transform_ffs = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
video_transforms.RandomHorizontalFlipVideo(),
video_transforms.UCFCenterCropVideo(args.image_size),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return FaceForensics(args, transform=transform_ffs, temporal_sample=temporal_sample)
elif args.dataset == 'ffs_img':
transform_ffs = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
video_transforms.RandomHorizontalFlipVideo(),
video_transforms.UCFCenterCropVideo(args.image_size),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return FaceForensicsImages(args, transform=transform_ffs, temporal_sample=temporal_sample)
elif args.dataset == 'ucf101':
transform_ucf101 = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
video_transforms.RandomHorizontalFlipVideo(),
video_transforms.UCFCenterCropVideo(args.image_size),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return UCF101(args, transform=transform_ucf101, temporal_sample=temporal_sample)
elif args.dataset == 'ucf101_img':
transform_ucf101 = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
video_transforms.RandomHorizontalFlipVideo(),
video_transforms.UCFCenterCropVideo(args.image_size),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return UCF101Images(args, transform=transform_ucf101, temporal_sample=temporal_sample)
elif args.dataset == 'taichi':
transform_taichi = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
video_transforms.RandomHorizontalFlipVideo(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return Taichi(args, transform=transform_taichi, temporal_sample=temporal_sample)
elif args.dataset == 'taichi_img':
transform_taichi = transforms.Compose([
video_transforms.ToTensorVideo(), # TCHW
video_transforms.RandomHorizontalFlipVideo(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return TaichiImages(args, transform=transform_taichi, temporal_sample=temporal_sample)
elif args.dataset == 'sky':
transform_sky = transforms.Compose([
video_transforms.ToTensorVideo(),
video_transforms.CenterCropResizeVideo(args.image_size),
# video_transforms.RandomHorizontalFlipVideo(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return Sky(args, transform=transform_sky, temporal_sample=temporal_sample)
elif args.dataset == 'sky_img':
transform_sky = transforms.Compose([
video_transforms.ToTensorVideo(),
video_transforms.CenterCropResizeVideo(args.image_size),
# video_transforms.RandomHorizontalFlipVideo(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
return SkyImages(args, transform=transform_sky, temporal_sample=temporal_sample)
else:
raise NotImplementedError(args.dataset) |