paulpanwang/PartCrafter

1import os2import json3import argparse4import time5from tqdm import tqdm6 7if __name__ == '__main__':8    parser = argparse.ArgumentParser()9    parser.add_argument('--input', type=str, default='assets/objects')10    parser.add_argument('--output', type=str, default='preprocessed_data')11    args = parser.parse_args()12 13    input_path = args.input14    output_path = args.output15 16    assert os.path.exists(input_path), f'{input_path} does not exist'17 18    if not os.path.exists(output_path):19        os.makedirs(output_path)20 21    for mesh_name in tqdm(os.listdir(input_path)):22        mesh_path = os.path.join(input_path, mesh_name)23        # 1. Sample points from mesh surface24        os.system(f"python datasets/preprocess/mesh_to_point.py --input {mesh_path} --output {output_path}")25        # 2. Render images26        os.system(f"python datasets/preprocess/render.py --input {mesh_path} --output {output_path}")27        # 3. Remove background for rendered images and resize to 90%28        export_mesh_folder = os.path.join(output_path, mesh_name.replace('.glb', ''))29        export_rendering_path = os.path.join(export_mesh_folder, 'rendering.png')30        os.system(f"python datasets/preprocess/rmbg.py --input {export_rendering_path} --output {output_path}")31        # 4. (Optional) Calculate IoU32        os.system(f"python datasets/preprocess/calculate_iou.py --input {mesh_path} --output {output_path}")33        time.sleep(1)34    35    # generate configs36    configs = []37    for mesh_name in tqdm(os.listdir(input_path)):38        mesh_path = os.path.join(output_path, mesh_name.replace('.glb', ''))39        num_parts_path = os.path.join(mesh_path, 'num_parts.json')40        surface_path = os.path.join(mesh_path, 'points.npy')41        image_path = os.path.join(mesh_path, 'rendering_rmbg.png')42        iou_path = os.path.join(mesh_path, 'iou.json')43        config = {44            "file": mesh_name,45            "num_parts": 0,46            "valid": False,47            "mesh_path": os.path.join(input_path, mesh_name),48            "surface_path": None,49            "image_path": None,50            "iou_mean": 0.0,51            "iou_max": 0.052        }53        try:54            config["num_parts"] = json.load(open(num_parts_path))['num_parts']55            iou_config = json.load(open(iou_path))56            config['iou_mean'] = iou_config['iou_mean']57            config['iou_max'] = iou_config['iou_max']58            assert os.path.exists(surface_path)59            config['surface_path'] = surface_path60            assert os.path.exists(image_path)61            config['image_path'] = image_path62            config['valid'] = True63            configs.append(config)64        except:65            continue66    67    configs_path = os.path.join(output_path, 'object_part_configs.json')68    json.dump(configs, open(configs_path, 'w'), indent=4)