multi-data

opensora/dataset/t2v_datasets.py

@@ -23,24 +23,23 @@ def random_video_noise(t, c, h, w):
 
 class T2V_dataset(Dataset):
     def __init__(self, args, transform, temporal_sample, tokenizer):
-        self.video_folder = args.video_folder
+        self.image_data = args.image_data
+        self.video_data = args.video_data
         self.num_frames = args.num_frames
         self.transform = transform
         self.temporal_sample = temporal_sample
         self.tokenizer = tokenizer
         self.model_max_length = args.model_max_length
-        self.video_folder = args.video_folder
         self.v_decoder = DecordInit()
 
-        with open(args.video_data_path, 'r') as f:
-            self.vid_cap_list = json.load(f)
+        self.vid_cap_list = self.get_vid_cap_list()
+        
         self.use_image_num = args.use_image_num
         self.use_img_from_vid = args.use_img_from_vid
         if self.use_image_num != 0 and not self.use_img_from_vid:
-            self.image_folder = args.image_folder
-            self.image_data_path = args.image_data_path
             self.img_cap_list = self.get_img_cap_list()
 
+
     def __len__(self):
         return len(self.vid_cap_list)
 
@@ -65,7 +64,7 @@ class T2V_dataset(Dataset):
         # input_ids = torch.ones(1, 120).to(torch.long).squeeze(0)
         # cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0)
         
-        video_path = opj(self.video_folder, self.vid_cap_list[idx]['path'])
+        video_path = self.vid_cap_list[idx]['path']
         video = self.decord_read(video_path, self.vid_cap_list[idx]['frame_idx'])
         video = self.transform(video)  # T C H W -> T C H W
         video = video.transpose(0, 1)  # T C H W -> C T H W
@@ -97,7 +96,7 @@ class T2V_dataset(Dataset):
         idx = idx % len(self.img_cap_list)  # out of range
         image_data = self.img_cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
         
-        image = [Image.open(os.path.join(self.image_folder, i['path'])).convert('RGB') for i in image_data] # num_img [h, w, c]
+        image = [Image.open(i['path']).convert('RGB') for i in image_data] # num_img [h, w, c]
         image = [torch.from_numpy(np.array(i)) for i in image] # num_img [h, w, c]
         image = [rearrange(i, 'h w c -> c h w').unsqueeze(0) for i in image] # num_img [1 c h w]
         image = [self.transform(i) for i in image]  # num_img [1 C H W] -> num_img [1 C H W]
@@ -152,8 +151,26 @@ class T2V_dataset(Dataset):
         video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
         return video_data
 
+
+    def get_vid_cap_list(self):
+        vid_cap_lists = []
+        with open(self.video_data, 'r') as f:
+            folder_anno = [i.strip().split(',') for i in f.readlines() if len(i.strip()) > 0]
+            for folder, anno in folder_anno:
+                vid_cap_list = json.load(anno)
+                for i in range(len(vid_cap_list)):
+                    vid_cap_list[i]['path'] = opj(folder, vid_cap_list[i]['path'])
+                vid_cap_lists += vid_cap_list
+        return vid_cap_lists
+
     def get_img_cap_list(self):
-        with open(self.image_data_path, 'r') as f:
-            image_data = json.load(f)
-        image_data = [image_data[i: i+self.use_image_num] for i in range(0, len(image_data), self.use_image_num)]
-        return image_data[:-1]  # drop last to avoid error length
+        img_cap_lists = []
+        with open(self.image_data, 'r') as f:
+            folder_anno = [i.strip().split(',') for i in f.readlines() if len(i.strip()) > 0]
+            for folder, anno in folder_anno:
+                img_cap_list = json.load(anno)
+                for i in range(len(img_cap_list)):
+                    img_cap_list[i]['path'] = opj(folder, img_cap_list[i]['path'])
+                img_cap_lists += img_cap_list
+        img_cap_lists = [img_cap_lists[i: i+self.use_image_num] for i in range(0, len(img_cap_lists), self.use_image_num)]
+        return img_cap_lists[:-1]  # drop last to avoid error length

opensora/train/train_t2v.py

@@ -604,10 +604,8 @@ def main(args):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--dataset", type=str, required=True)
-    parser.add_argument("--video_data_path", type=str, required=True)
-    parser.add_argument("--video_folder", type=str, default='')
-    parser.add_argument("--image_data_path", type=str, default='')
-    parser.add_argument("--image_folder", type=str, default='')
+    parser.add_argument("--video_data", type=str, required='')
+    parser.add_argument("--image_data", type=str, default='')
     parser.add_argument("--sample_rate", type=int, default=1)
     parser.add_argument("--num_frames", type=int, default=17)
     parser.add_argument("--max_image_size", type=int, default=512)

scripts/train_data/image_data.txt

@@ -0,0 +1,2 @@
+/remote-home1/dataset/picture,/remote-home1/yeyang/gan3d/Open-Sora-Plan/image_114054.json
+/remote-home1/dataset/picture,/remote-home1/yeyang/gan3d/Open-Sora-Plan/image_114054.json

scripts/train_data/video_data.txt

@@ -0,0 +1,2 @@
+/remote-home1/dataset/all_mixkit,/remote-home1/yeyang/gan3d/Open-Sora-Plan/video_mixkit_65f_54735.json
+/remote-home1/dataset/all_mixkit,/remote-home1/yeyang/gan3d/Open-Sora-Plan/video_mixkit_65f_54735.json