update

mmhuman3d/data/data_converters/ubody.py

@@ -223,19 +223,19 @@ def convert_by_mode(self, dataset_path: str, out_path: str,
         #         test_vids += vid_ps[int((len(vid_ps)+1)*.70):]
         # processed_vids += vid_ps
         # processed_vids = list(dict.fromkeys(processed_vids))
-        for scene in scene_split:
-            if not '_' in scene:
-                continue
-            vid_ps = [vid_p for vid_p in vid_ps_all if scene in vid_p]
-            test_vids += vid_ps
-        test_vids = list(dict.fromkeys(test_vids))
-        train_vids = [vid_p for vid_p in vid_ps_all if vid_p not in test_vids]
-        # pdb.set_trace()
+
+
+        # for scene in scene_split:
+        #     if not '_' in scene:
+        #         continue
+        #     vid_ps = [vid_p for vid_p in vid_ps_all if scene in vid_p]
+        #     test_vids += vid_ps
+        # test_vids = list(dict.fromkeys(test_vids))
+        # train_vids = [vid_p for vid_p in vid_ps_all if vid_p not in test_vids]
 
         # vid_ps = vid_ps[:1]
 
-        for batch, vid_ps_batch in zip(['test', 'train'],
-                                       [test_vids, train_vids]):
+        for batch in ['test', 'train']:
 
             # num_proc = 4
             # with Pool(num_proc) as p:
@@ -277,17 +277,16 @@ def convert_by_mode(self, dataset_path: str, out_path: str,
                 # load seq smplx annotation
                 with open(os.path.join(anno_folder, 'smplx_annotation.json')) as f:
                     smplx_param = json.load(f)
-                vid_ps_scene = [v for v in vid_ps_batch if vscene in v]
                 # for vid_p in tqdm(vid_ps_scene, desc=f'Preprocessing scene {bid+1}/{len(vid_batches)}', 
                 #                   position=0, leave=False):
                 # process_vid(vid_p, smplx_model, anno_param, smplx_param)
 
-                process_vid_COCO(smplx_model, vscene, vid_ps_batch, db, smplx_param, dst)
+                process_vid_COCO(smplx_model, vscene, scene_split, batch, db, smplx_param, dst)
 
                 # pdb.set_trace()
 
             processed_npzs = glob.glob(os.path.join(dst, '*.npz')) 
-            seed, size = '231027', '99'
+            seed, size = '240409', '99'
 
             random.seed(int(seed))
 
@@ -439,7 +438,7 @@ def convert_by_mode(self, dataset_path: str, out_path: str,
                       time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))
 
                 # store
-                human_data['config'] = f'egobody_{mode}'
+                human_data['config'] = f'ubody_{mode}'
                 human_data['misc'] = self.misc_config
 
                 human_data.compress_keypoints_by_mask()

tools/preprocess/ubody_preprocess.py

@@ -322,7 +322,7 @@ def process_vid(vid, smplx_model, anno_param, smplx_param):
     #     meta_['principal_point'].append(principal_point)
 
 
-def process_vid_COCO(smplx_model, vscene, vid_ps_batch, db, smplx_param, dst):
+def process_vid_COCO(smplx_model, vscene, scene_split, batch, db, smplx_param, dst):
     # vid = args.vid_p
     device = torch.device(
         'cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -373,9 +373,10 @@ def process_vid_COCO(smplx_model, vscene, vid_ps_batch, db, smplx_param, dst):
     aids = list(db.anns.keys())
 
     # prepare vid list 
-    vid_list = [vid_p.split('/')[-2] for vid_p in vid_ps_batch]
+    # vid_list = [vid_p.split('/')[-2] for vid_p in vid_ps_batch]
     dataset_path = os.path.sep.join(dst.split('/')[:-2])
 
+    valid_count = 0
     for bid in range(blocks):
 
         param_dict = {}
@@ -403,10 +404,26 @@ def process_vid_COCO(smplx_model, vscene, vid_ps_batch, db, smplx_param, dst):
             image_info = db.loadImgs(ann['image_id'])[0]
             video_name = image_info['file_name'].split('/')[-2]
 
-            if not video_name in vid_list:
-                continue
+            if image_info['file_name'].startswith('/'):
+                file_name = image_info['file_name'][1:]  # [1:] means delete '/'
+            else:
+                file_name = image_info['file_name']
+
+            video_name = file_name.split('/')[-2]
+            if 'Trim' in video_name:
+                video_name = video_name.split('_Trim')[0]
+
+            if batch == 'train':
+                if video_name in scene_split:
+                    continue
+            elif batch == 'test':
+                if video_name not in scene_split:
+                    continue       
+
             imgp = os.path.join(dataset_path, 'images', 
                                 vscene, image_info['file_name'])
+            if ann['iscrowd'] or (ann['num_keypoints']==0): continue
+            if ann['valid_label'] == 0: continue
             if not os.path.exists(imgp): 
                 continue
             if str(aid) not in smplx_param: 
@@ -518,11 +535,13 @@ def process_vid_COCO(smplx_model, vscene, vid_ps_batch, db, smplx_param, dst):
             for key in ['iscrowd', 'num_keypoints', 'valid_label', 'lefthand_valid',
                 'righthand_valid', 'face_valid']:
                 param_dict[key].append(kp_param[key])
+            valid_count += 1
             # pdb.set_trace()
 
         os.makedirs(os.path.dirname(preprocess_file), exist_ok=True)
         np.savez(preprocess_file, **param_dict)
 
+    print(f'{vscene}: Valid count: {valid_count}')
 
 
 if __name__ == '__main__':