From 8dc49b8a6b9a737bbfad803b8ffff7b2ce4aff18 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Mon, 29 Apr 2024 19:01:10 +0800 Subject: [PATCH] multi-data --- opensora/dataset/t2v_datasets.py | 41 ++++++++++++++++++++++--------- opensora/train/train_t2v.py | 6 ++--- scripts/train_data/image_data.txt | 2 ++ scripts/train_data/video_data.txt | 2 ++ 4 files changed, 35 insertions(+), 16 deletions(-) create mode 100644 scripts/train_data/image_data.txt create mode 100644 scripts/train_data/video_data.txt diff --git a/opensora/dataset/t2v_datasets.py b/opensora/dataset/t2v_datasets.py index c548a453c..818becdfa 100644 --- a/opensora/dataset/t2v_datasets.py +++ b/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 \ No newline at end of file + 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 \ No newline at end of file diff --git a/opensora/train/train_t2v.py b/opensora/train/train_t2v.py index b18828f6b..095472d3f 100644 --- a/opensora/train/train_t2v.py +++ b/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) diff --git a/scripts/train_data/image_data.txt b/scripts/train_data/image_data.txt new file mode 100644 index 000000000..936f647bb --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/train_data/video_data.txt b/scripts/train_data/video_data.txt new file mode 100644 index 000000000..2419f6d0a --- /dev/null +++ b/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 \ No newline at end of file