add occlusion splitting

compute_label.py

@@ -12,11 +12,13 @@ if __name__ == "__main__":
     parser.add_argument('--path_dataset', type=str, required=True)
     parser.add_argument('--target_folder', type=str, required=True)
     parser.add_argument('--path_DF_data', type=str, required=True)
+    parser.add_argument('--occ_data', type=str, required=True)
     # Parse the argument
     args = parser.parse_args()
 
     target_folder = args.target_folder
     path_data_df = args.path_DF_data
+    occ_data = args.occ_data
 
     #objects = ["apple2", "banana1", "pear2"]
     objects = ["apple2", "apricot", "banana1", "kiwi1", "lemon2", "orange2", "peach1", "pear2"]
@@ -29,7 +31,7 @@ if __name__ == "__main__":
     #path_data = "/gpfsscratch/rech/uli/ubn15wo/dataset1/s2rg/Fruits_all_medium/GUIMOD_low/Generated_Worlds_Training/orange2/Pose_transformed/"
     #path_data = f"/gpfsscratch/rech/uli/ubn15wo/dataset1/s2rg/Fruits_all_medium/GUIMOD_low/{target_folder}/"
 
-    for split in ["_Training","_Evaluating"]:
+    for split in ["_Training", "_Evaluating", "_Testing"]:
         path_data = f"{args.path_dataset}/{target_folder}{split}"
 
         if not os.path.exists(path_data_df) : 
@@ -59,22 +61,22 @@ if __name__ == "__main__":
         for obj in objects:
             file = path_data + "/" + obj
             #path_gt_obj = f"{obj_ids[obj]}/gt_new.yml"
-            path_gt_obj = f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}{split}.yml"
+            path_gt_obj = f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}{split}{occ_data}.yml"
             #path_bbox = "GUIMOD_low/Bbox/" + obj
             path_bbox = f"{path_data}/{obj}/Bbox/"
 
             object_dicts = {}
             print(file)
             path_poses = file + "/Pose_transformed"
-            print("poses ", path_poses)
-            print("list : ", os.listdir(path_poses))
+            #print("poses ", path_poses)
+            #print("list : ", os.listdir(path_poses))
             list_poses = os.listdir(path_poses)
 
             print("-------------------------------------------------------------------------------------------------------")
             print("-------------------------------------------------------------------------------------------------------")
             
             if split=="_Training" : 
-                with open(f"{path_data_df}/0{obj_ids[obj]}/train.txt", 'w') as f:
+                with open(f"{path_data_df}/0{obj_ids[obj]}/train{occ_data}.txt", 'w') as f:
                     for filename2 in os.listdir(path_poses):
                         id_file2=filename2.split('.')[0]
                         print(id_file2+".png")
@@ -82,14 +84,22 @@ if __name__ == "__main__":
                         f.write('\n')
 
             elif split=="_Evaluating" : 
-                with open(f"{path_data_df}/0{obj_ids[obj]}/test.txt", 'w') as f:
+                with open(f"{path_data_df}/0{obj_ids[obj]}/test{occ_data}.txt", 'w') as f:
                     for filename2 in os.listdir(path_poses):
                         id_file2=filename2.split('.')[0]
                         print(id_file2+".png")
                         f.write(f'{id_file2}.png')
                         f.write('\n')
+            elif split=="_Testing" : 
+                with open(f"{path_data_df}/0{obj_ids[obj]}/test_true{occ_data}.txt", 'w') as f:
+                    for filename2 in os.listdir(path_poses):
+                        id_file2=filename2.split('.')[0]
+                        print(id_file2+".png")
+                        f.write(f'{id_file2}.png')
+                        f.write('\n')
+
             else : 
-                raise("not in traing or evaluating")
+                raise("not in traing or evaluating or testing")
 
             with open(path_gt_obj, 'w') as yaml_file:
                 for filename in os.listdir(path_poses):
@@ -123,18 +133,27 @@ if __name__ == "__main__":
     for obj in objects:
 
         # Reading data from file1
-        with open(f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}_Training.yml") as fp:
+        with open(f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}_Training{occ_data}.yml") as fp:
             data_train = fp.read()
         
         # Reading data from file2
-        with open(f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}_Evaluating.yml") as fp:
+        with open(f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}_Evaluating{occ_data}.yml") as fp:
             data_eval = fp.read()
+
+        # Reading data from file2
+        with open(f"{path_data_df}/0{obj_ids[obj]}/gt_{target_folder}_Testing{occ_data}.yml") as fp:
+            data_test = fp.read()
+
         
         # Merging 2 files
         # To add the data of file2
         # from next line
         data_gt = data_train
         data_gt += data_eval
-        
-        with open (f"{path_data_df}/0{obj_ids[obj]}/gt.yml", 'w') as fp:
-            fp.write(data_gt)
\ No newline at end of file
+        data_gt_true = data_gt + data_test
+        with open (f"{path_data_df}/0{obj_ids[obj]}/gt{occ_data}.yml", 'w') as fp:
+            fp.write(data_gt)
+
+        with open (f"{path_data_df}/0{obj_ids[obj]}/gt_true{occ_data}.yml", 'w') as fp:
+            fp.write(data_gt_true)
+