diff --git a/main.py b/main.py new file mode 100644 index 0000000000000000000000000000000000000000..90c471b4d07237ef0dd9c8aca06a528d74d2b2c1 --- /dev/null +++ b/main.py @@ -0,0 +1,63 @@ +import os +import numpy as np +from prepare_data import reform_data +from pose import transform_pose +from bbox_2d import generate_2d_bbox +from instance_mask import generate_instance_mask +from fps_alg import generate_fps +from bbox_3d import generate_3d_bbox +import shutil + + +def generate_folders(name): + is_exist = os.path.exists(name) + if not is_exist: + os.mkdir(name) + + folders = ["RGB", "Depth", "Mask", "Pose", "Bbox_2d", "Instance_Mask", "Models", "Pose_transformed", "Bbox", "FPS"] + for f in folders: + is_exist = os.path.exists(f"{name}/{f}") + if not is_exist: + if f not in ["Instance_Mask", "Models", "Pose_transformed", "Bbox", "FPS"]: + os.mkdir(f"{name}/{f}") + else: + os.makedirs(f"{name}/{f}/Banana") + os.makedirs(f"{name}/{f}/Orange") + os.makedirs(f"{name}/{f}/Pear") + + +if __name__ == '__main__': + dataset_src = "/media/mahmoud/E/Fruits_easy/data" + choice = "low" + + data_options = {"high": "ground_truth_rgb", + "low": "ground_truth_depth"} + + dataset_type = data_options[choice] + + dataset_name = f"GUIMOD_{choice}" + generate_folders(dataset_name) + + list_obj = ["Banana", "Orange", "Pear"] + for obj in list_obj: + src_bbox = f"Models/{obj}/{obj.lower()}.ply" + dst_bbox = f"{dataset_name}/Models/{obj}/{obj.lower()}.ply" + shutil.copy(src_bbox, dst_bbox) + + if choice == 'high': + camera = np.matrix([[1386.4138492513919, 0.0, 960.5], + [0.0, 1386.4138492513919, 540.5], + [0.0, 0.0, 1.0]]) + else: + camera = np.matrix([[1086.5054444841007, 0.0, 640.5], + [0.0, 1086.5054444841007, 360.5], + [0.0, 0.0, 1.0]]) + + np.savetxt(f'{dataset_name}/camera_{choice}.txt', camera) + + reform_data(dataset_src, dataset_name, dataset_type) + transform_pose(dataset_name) + generate_2d_bbox(dataset_name) + generate_instance_mask(dataset_name) + generate_fps(dataset_name, camera, False) + generate_3d_bbox(dataset_name)