Spaces:
Paused
Paused
File size: 6,340 Bytes
96257b2 |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
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 .buv_datasets import BUV
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
# from .echonet_datasets import EchoNetDynamic
from .echonet_datasets_noise_dual import EchoNetDynamic
from .echonet_datasets_inversed_noise import EchoNetDynamic_Noise
from .echonet_datasets_inversed_noise_timestep import EchoNetDynamic_Noise_Timestep
from .coche_datasets import COCHE
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 == 'coche':
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 COCHE(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 == 'buv':
transform_buv = 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 BUV(args, transform=transform_buv, 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)
elif args.dataset == 'echonet':
transform_echo = 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 EchoNetDynamic(args, transform=transform_echo, temporal_sample=temporal_sample)
elif args.dataset == 'echonet_noise':
transform_echo = None
return EchoNetDynamic_Noise(args, transform=transform_echo, temporal_sample=temporal_sample)
elif args.dataset == 'echonet_noise_timestep':
transform_echo = None
return EchoNetDynamic_Noise_Timestep(args, transform=transform_echo, temporal_sample=temporal_sample)
else:
raise NotImplementedError(args.dataset) |