drema_project/create_simulation.py at main · berkecyln/drema_project · 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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import os.path
import time
import hydra

from omegaconf import DictConfig, OmegaConf

from drema.environment.assets.object import Object
from drema.r2s_builder.assets_extractor import AssetsManager
from drema.r2s_builder.gaussians_optimizers.depth_optimizer import DepthTrainer
from drema.r2s_builder.gaussians_optimizers.surf_depth_optimizer import SurfDepthTrainer
from drema.r2s_builder.gaussians_optimizers.surf_optimizer import SurfTrainer
from drema.r2s_builder.gaussians_optimizers.base_optimizer import BaseTrainer
from drema.utils.coppelia_utils import read_labels

def prepare_labels(path):
    labels = read_labels(path, filter_labels=False)

    filter_names = ["DefaultCamera", "ResizableFloor", "workspace", "Wall"]
    # remove labels containing the filter names
    labels = {k: v for k, v in labels.items() if not any(name in k for name in filter_names)}

    counter = 1
    print(labels)
    Panda_labels = {}
    Object_labels = {}
    Table_labels = {}
    # iterate over the labels
    for label_name, value in labels.items():
        # check if the label contain the Panda string
        if "Panda" in label_name:
            label_name = label_name.split("_")[1]
            if label_name == "gripper":
                label_name = "link8"
            elif label_name == "leftfinger":
                label_name = "link10"
            elif label_name == "rightfinger":
                label_name = "link9"
            Panda_labels[label_name] = value
        elif label_name in ["table", "diningTable_visible", "surface"]:
            Table_labels[label_name] = value
        else:
            # check if label_name is in keys
            if label_name in Table_labels.keys():
                label_name = label_name + "_"+str(counter)
                counter += 1
            Object_labels[label_name] = value

    #ADDED ==
    print("Panda Labels: ", Panda_labels)
    print("Object Labels: ", Object_labels)
    print("Table Labels: ", Table_labels)
    #========

    return Panda_labels, Object_labels, Table_labels


@hydra.main(version_base=None, config_path="configs", config_name="config_real")
def main(cfg: DictConfig) -> None:
    print("Building the simulation")

    source_path = cfg.data.source_path
    assets_path = cfg.data.assets_path

    dataset = cfg.training.model
    pipeline = cfg.training.pipeline
    optimization = cfg.training.optimization

    assets = cfg.training.assets

    gaussians_iterations = assets.gaussians_iterations
    mesh_iterations = optimization.iterations

    if assets.use_original_guassians:
        if assets.use_depth:
            trainer = DepthTrainer
        else:
            trainer = BaseTrainer
    else:
        if assets.use_depth:
            trainer = SurfDepthTrainer
        else:
            trainer = SurfTrainer

    print(f"Using trainer: {trainer.__name__}")

    experiment_name = f"Experiment_{time.strftime('%Y%m%d-%H%M%S')}"

    assets_manager = AssetsManager(source_path, assets_path, trainer, dataset, optimization, pipeline, gaussians_iterations, mesh_iterations, experiment_name, assets.filter_mesh_objects, assets.fill_mesh_holes, assets.mesh_method, assets.depth_dir)

    # Load the data
    assets_manager.load_data()

    # Extract Environment
    if assets.extract_gaussians_environment:
        print("================Extracting Environment================")
        assets_manager.extract_environment(assets.extract_mesh_environment)
        print("================Extracting Environment Done================")

    # get labels
    print("================Preparing Labels================")
    panda_labels, object_labels, table_labels = prepare_labels(os.path.join(source_path, "labels.txt"))
    print("================Preparing Labels Done================")

    # if it needs to extract the table
    if assets.extract_table:
        # extract the table
        # TODO: extract table even without depth
        table_labels = list(table_labels.values())[0]
        print("================Table Extraction================")
        assets_manager.extract_table(table_labels)
        print("================Table Extraction Done================")
    else:
        # load the table
        assets_manager.load_table()

    if assets.extract_gaussians_objects:
        for label_name, value in object_labels.items():
            # extract the object
            print(f"================Extracting Object: {label_name}================")
            assets_manager.extract_asset(value, extract_mesh=assets.extract_mesh_objects, extract_urdf=assets.extract_urdf_objects)
            print(f"================Extracting Object Done: {label_name}================")

    if assets.extract_gaussians_robot:
        for label_name, value in panda_labels.items():
            # extract the object
            print(f"================Extracting Robot Part: {label_name}================")
            assets_manager.extract_asset(value, extract_mesh=False, extract_urdf=False)
            print(f"================Extracting Robot Part Done: {label_name}================")

            # move the output to the correct folder
            src = os.path.join(source_path, "output", "objects_ply", str(value) + ".ply")
            dst = os.path.join(source_path, "output", "robot", label_name + ".ply")
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            os.rename(src, dst)

    if assets.filter_objects_gaussians_environment:
        print("================Filtering Environment================") #ADDED
        assets_manager.filter_environment()
        print("================Filtering Environment Done================") #ADDED


if __name__ == "__main__":
    main()