feature dataset

Former-commit-id: 9683caae8648735552270b871de22fc540ba2b55

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(
             [

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

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

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