PartCrafter/datasets/preprocess/preprocess.py at main · AveryRPeterson/PartCrafter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import os
import json
import argparse
import time
from tqdm import tqdm

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default='assets/objects')
    parser.add_argument('--output', type=str, default='preprocessed_data')
    args = parser.parse_args()

    input_path = args.input
    output_path = args.output

    assert os.path.exists(input_path), f'{input_path} does not exist'

    if not os.path.exists(output_path):
        os.makedirs(output_path)

    for mesh_name in tqdm(os.listdir(input_path)):
        mesh_path = os.path.join(input_path, mesh_name)
        # 1. Sample points from mesh surface
        os.system(f"python datasets/preprocess/mesh_to_point.py --input {mesh_path} --output {output_path}")
        # 2. Render images
        os.system(f"python datasets/preprocess/render.py --input {mesh_path} --output {output_path}")
        # 3. Remove background for rendered images and resize to 90%
        export_mesh_folder = os.path.join(output_path, mesh_name.replace('.glb', ''))
        export_rendering_path = os.path.join(export_mesh_folder, 'rendering.png')
        os.system(f"python datasets/preprocess/rmbg.py --input {export_rendering_path} --output {output_path}")
        # 4. (Optional) Calculate IoU
        os.system(f"python datasets/preprocess/calculate_iou.py --input {mesh_path} --output {output_path}")
        time.sleep(1)

    # generate configs
    configs = []
    for mesh_name in tqdm(os.listdir(input_path)):
        mesh_path = os.path.join(output_path, mesh_name.replace('.glb', ''))
        num_parts_path = os.path.join(mesh_path, 'num_parts.json')
        surface_path = os.path.join(mesh_path, 'points.npy')
        image_path = os.path.join(mesh_path, 'rendering_rmbg.png')
        iou_path = os.path.join(mesh_path, 'iou.json')
        config = {
            "file": mesh_name,
            "num_parts": 0,
            "valid": False,
            "mesh_path": os.path.join(input_path, mesh_name),
            "surface_path": None,
            "image_path": None,
            "iou_mean": 0.0,
            "iou_max": 0.0
        }
        try:
            config["num_parts"] = json.load(open(num_parts_path))['num_parts']
            iou_config = json.load(open(iou_path))
            config['iou_mean'] = iou_config['iou_mean']
            config['iou_max'] = iou_config['iou_max']
            assert os.path.exists(surface_path)
            config['surface_path'] = surface_path
            assert os.path.exists(image_path)
            config['image_path'] = image_path
            config['valid'] = True
            configs.append(config)
        except:
            continue

    configs_path = os.path.join(output_path, 'object_part_configs.json')
    json.dump(configs, open(configs_path, 'w'), indent=4)