compute label for yolo, gt.yml, training.txt, testing.txt for Densefusion

compute_label.py

+import numpy as np
+import yaml
+import os
+import argparse
+
+
+if __name__ == "__main__":
+    # Create the parser
+    parser = argparse.ArgumentParser()
+    # Add an argument
+    parser.add_argument('--path_dataset', type=str, required=True)
+    parser.add_argument('--target_folder', type=str, required=True)
+    # Parse the argument
+    args = parser.parse_args()
+
+    target_folder = args.target_folder
+
+    #objects = ["apple2", "banana1", "pear2"]
+    objects = ["apple2", "apricot", "banana1", "kiwi1", "lemon2", "orange2", "peach1", "pear2"]
+    #obj_ids = {"pear2": 17, "orange2": 18}
+
+    obj_ids = {"apple2": 1, "apricot": 2, "banana1": 3, "kiwi1": 4, "lemon2": 5, "orange2": 6, "peach1": 7, "pear2": 8}
+    #obj_ids = {"apple2": 1, "banana1": 3,  "pear2": 8}
+    
+    #path_data = "GUIMOD_low/Pose_transformed"
+    #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}/"
+    path_data = f"{args.path_dataset}/{target_folder}"
+
+    for obj in objects:
+        file = path_data + "/" + obj
+        #path_gt_obj = f"{obj_ids[obj]}/gt_new.yml"
+        path_gt_obj = f"/gpfsscratch/rech/uli/ubn15wo/DenseFusion/datasets/linemod/Linemod_preprocessed/data/0{obj_ids[obj]}/gt_{target_folder}.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))
+        list_poses = os.listdir(path_poses)
+
+        print("-------------------------------------------------------------------------------------------------------")
+        print("-------------------------------------------------------------------------------------------------------")
+        with open(f"/gpfsscratch/rech/uli/ubn15wo/DenseFusion/datasets/linemod/Linemod_preprocessed/data/0{obj_ids[obj]}/split_file_{target_folder}.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')
+
+        with open(path_gt_obj, 'w') as yaml_file:
+            for filename in os.listdir(path_poses):
+            #for i in range(num_rgb):
+                data_file = path_poses + "/" + filename
+                print("filename : ", filename)
+                #data_file = file + "/" + str(i) + ".npy"
+                #print(list_RGB[4].split('.')[0])
+                id_file=filename.split('.')[0]
+                print(id_file)
+                print(filename)
+                #final_path_bbox = path_bbox + "/" + str(i) + ".txt"
+                final_path_bbox = path_bbox + "/" + id_file + ".txt"
+                print(final_path_bbox)
+                bbox_file = open(final_path_bbox, "r")
+                line = bbox_file.readline()
+                bbox_string = line.split(" ")
+                bbox = [float(val) for val in bbox_string]
+                bbox2 = [int(bbox[0] * 0.5), int(bbox[1] * (2 / 3)), int(bbox[2] * 0.5), int(bbox[3] * (2 / 3))]
+                # print(bbox2)
+                bbox_file.close()
+
+                data = np.load(data_file)
+                cam_R_m2c = np.concatenate((data[0, 0:3], data[1, 0:3], data[2, 0:3]), axis=0).tolist()
+                cam_t_m2c = np.array((float(data[0, 3])*1000.0, float(data[1, 3])*1000.0, float(data[2, 3]*1000.0))).tolist()
+                object_dicts[id_file] = [{"cam_R_m2c": cam_R_m2c, "cam_t_m2c": cam_t_m2c, "obj_bb": bbox2, "obj_id": obj_ids[obj]}]
+            yaml.dump(object_dicts, yaml_file, default_flow_style=None)