From b07538bf8a5ddce73e494725b23d413cf140ce8f Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Fri, 8 Mar 2024 17:24:23 +0800 Subject: [PATCH] feature dataset Former-commit-id: 9683caae8648735552270b871de22fc540ba2b55 --- opensora/dataset/__init__.py | 9 ++- opensora/dataset/extract_feature_dataset.py | 90 +++++++++++++++++++++ opensora/dataset/feature_datasets.py | 73 +++++++++++++++-- opensora/dataset/sky_datasets.py | 10 +-- 4 files changed, 167 insertions(+), 15 deletions(-) create mode 100644 opensora/dataset/extract_feature_dataset.py diff --git a/opensora/dataset/__init__.py b/opensora/dataset/__init__.py index f7d28977d..d0b13ff02 100644 --- a/opensora/dataset/__init__.py +++ b/opensora/dataset/__init__.py @@ -2,7 +2,7 @@ from torchvision.transforms import Compose from opensora.models.ae import vae, vqvae, videovae, videovqvae -from .feature_datasets import LandscopeFeatures +from .feature_datasets import LandscopeFeature, SkyFeature from torchvision import transforms from torchvision.transforms import Lambda @@ -39,8 +39,11 @@ def getdataset(args): temporal_sample = TemporalRandomCrop(args.num_frames * args.sample_rate) # 16 x norm_fun = ae_norm[args.ae] if args.dataset == 'landscope_feature': - temporal_sample = TemporalRandomCrop(args.num_frames) # 16 1 - return LandscopeFeatures(args, temporal_sample=temporal_sample) + transform = Compose([CenterCropResizeVideo(size=args.max_image_size)]) + return LandscopeFeature(args, transform=transform, temporal_sample=temporal_sample) + elif args.dataset == 'sky_feature': + transform = Compose([CenterCropResizeVideo(size=args.max_image_size)]) + return SkyFeature(args, transform=transform, temporal_sample=temporal_sample) elif args.dataset == 'ucf101': transform = Compose( [ diff --git a/opensora/dataset/extract_feature_dataset.py b/opensora/dataset/extract_feature_dataset.py new file mode 100644 index 000000000..c7bf51af4 --- /dev/null +++ b/opensora/dataset/extract_feature_dataset.py @@ -0,0 +1,90 @@ +import os +from glob import glob + +import numpy as np +import torch +import torchvision +from PIL import Image +from torch.utils.data import Dataset + +from opensora.utils.dataset_utils import DecordInit, is_image_file + + +class ExtractVideo2Feature(Dataset): + def __init__(self, args, transform): + self.data_path = args.data_path + self.sample_rate = args.sample_rate + self.transform = transform + self.v_decoder = DecordInit() + self.samples = list(glob(f'{self.data_path}')) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + video_path = self.samples[idx] + video = self.decord_read(video_path) + video = self.transform(video) # T C H W -> T C H W + return video, video_path + + def tv_read(self, path): + vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW') + total_frames = len(vframes) + frame_indice = list(range(total_frames)) + video = vframes[frame_indice] + return video + + def decord_read(self, path): + decord_vr = self.v_decoder(path) + total_frames = len(decord_vr) + frame_indice = list(range(total_frames)) + video_data = decord_vr.get_batch(frame_indice).asnumpy() + video_data = torch.from_numpy(video_data) + video_data = video_data.permute(0, 3, 1, 2) # (T, H, W, C) -> (T C H W) + return video_data + + + +class ExtractImage2Feature(Dataset): + def __init__(self, args, transform): + self.data_path = args.data_path + self.transform = transform + self.data_all = self.load_video_frames(self.data_path) + + def __len__(self): + return len(self.data_all) + + def __getitem__(self, index): + vframes = self.data_all[index] + total_frames = len(vframes) + + # Sampling video frames + select_video_frames = vframes + + video_frames = [] + for path in select_video_frames: + video_frame = torch.as_tensor(np.array(Image.open(path), dtype=np.uint8, copy=True)).unsqueeze(0) + video_frames.append(video_frame) + video_clip = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2) + video_clip = self.transform(video_clip) # T C H W + # video_clip = video_clip.transpose(0, 1) # T C H W -> C T H W + + return video_clip, select_video_frames + + + def load_video_frames(self, dataroot): + data_all = [] + frame_list = os.walk(dataroot) + for _, meta in enumerate(frame_list): + root = meta[0] + try: + frames = sorted(meta[2], key=lambda item: int(item.split('.')[0].split('_')[-1])) + except: + pass + # print(meta[0]) # root + # print(meta[2]) # files + frames = [os.path.join(root, item) for item in frames if is_image_file(item)] + if len(frames) > 0: + data_all.append(frames) + self.video_num = len(data_all) + return data_all \ No newline at end of file diff --git a/opensora/dataset/feature_datasets.py b/opensora/dataset/feature_datasets.py index 9a0af726b..ad41f6225 100644 --- a/opensora/dataset/feature_datasets.py +++ b/opensora/dataset/feature_datasets.py @@ -11,19 +11,20 @@ from tqdm import tqdm from opensora.dataset.transform import center_crop, RandomCropVideo -class LandscopeFeatures(data.Dataset): - def __init__(self, args, temporal_sample): +class LandscopeFeature(data.Dataset): + def __init__(self, args, transform, temporal_sample): self.args = args self.data_path = args.data_path self.max_image_size = args.latent_size self.temporal_sample = temporal_sample - self.randomcrop = RandomCropVideo(self.max_image_size) + self.transform = transform data_all = list(glob(self.data_path)) self.num_frames = self.args.num_frames + self.sample_rate = self.args.sample_rate print('Building dataset...') - self.data_all = [i for i in tqdm(data_all) if self.num_frames < np.load(i).shape[1]] + self.data_all = [i for i in tqdm(data_all) if self.num_frames * self.sample_rate < np.load(i).shape[1]] print(f'Total {len(self.data_all)} to train, origin dataset have {len(data_all)} videos.') def __getitem__(self, index): @@ -37,8 +38,8 @@ class LandscopeFeatures(data.Dataset): start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) assert end_frame_ind - start_frame_ind >= self.num_frames video_clip = vframes[start_frame_ind: end_frame_ind] # - video_clip = torch.from_numpy(video_clip) - video_clip = self.randomcrop(video_clip) + video_clip = torch.from_numpy(video_clip) # T C H W + video_clip = self.transform(video_clip) video_clip = video_clip # T C H W return video_clip, 1 except Exception as e: @@ -49,3 +50,63 @@ class LandscopeFeatures(data.Dataset): def __len__(self): return len(self.data_all) + +class SkyFeature(data.Dataset): + def __init__(self, args, transform, temporal_sample=None): + + self.args = args + self.data_path = args.data_path + self.transform = transform + self.temporal_sample = temporal_sample + self.num_frames = self.args.num_frames + self.sample_rate = self.args.sample_rate + self.data_all = self.load_video_frames(self.data_path) + + def __getitem__(self, index): + + try: + vframes = self.data_all[index] + total_frames = len(vframes) + + # Sampling video frames + start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) + assert end_frame_ind - start_frame_ind >= self.num_frames + frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, num=self.num_frames, + dtype=int) # start, stop, num=50 + + select_video_frames = vframes[frame_indice[0]: frame_indice[-1] + 1: self.sample_rate] + + video_frames = [] + for path in select_video_frames: + video_frame = torch.from_numpy(np.load(path)).unsqueeze(0) # 1 c h w + video_frames.append(video_frame) + video_clip = torch.cat(video_frames, dim=0) # T C H W + video_clip = self.transform(video_clip) # T C H W + video_clip = video_clip.transpose(0, 1) # T C H W -> C T H W + + return video_clip, 1 + except Exception as e: + print(f'Error with {e}', vframes) + return self.__getitem__(random.randint(0, self.__len__()-1)) + + def __len__(self): + return self.video_num + + def load_video_frames(self, dataroot): + data_all = [] + frame_list = os.walk(dataroot) + for _, meta in enumerate(frame_list): + root = meta[0] + try: + frames = sorted(meta[2], key=lambda item: int(item.split('.')[0].split('_')[-1])) + except: + pass + # print(meta[0]) # root + # print(meta[2]) # files + frames = [os.path.join(root, item) for item in frames if item.endswith('.npy')] + if len(frames) > max(0, + self.target_video_len * self.frame_interval): # need all > (16 * frame-interval) videos + # if len(frames) >= max(0, self.target_video_len): # need all > 16 frames videos + data_all.append(frames) + self.video_num = len(data_all) + return data_all \ No newline at end of file diff --git a/opensora/dataset/sky_datasets.py b/opensora/dataset/sky_datasets.py index 487998b51..f491b4466 100644 --- a/opensora/dataset/sky_datasets.py +++ b/opensora/dataset/sky_datasets.py @@ -7,10 +7,8 @@ import numpy as np from PIL import Image -IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG'] +from opensora.utils.dataset_utils import is_image_file -def is_image_file(filename): - return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) class Sky(data.Dataset): def __init__(self, configs, transform, temporal_sample=None, train=True): @@ -56,9 +54,9 @@ class Sky(data.Dataset): try: frames = sorted(meta[2], key=lambda item: int(item.split('.')[0].split('_')[-1])) except: - # pass - print(meta[0]) # root - print(meta[2]) # files + pass + # print(meta[0]) # root + # print(meta[2]) # files frames = [os.path.join(root, item) for item in frames if is_image_file(item)] if len(frames) > max(0, self.target_video_len * self.frame_interval): # need all > (16 * frame-interval) videos # if len(frames) >= max(0, self.target_video_len): # need all > 16 frames videos