OpenReal2Sim/openreal2sim/motion/motion_manager.py at main · PointsCoder/OpenReal2Sim · 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
from pathlib import Path
import sys
import pickle
import yaml
import argparse
import json
sys.path.append(str(Path.cwd() / "openreal2sim" / "motion" / "modules"))
sys.path.append(str(Path.cwd() / "openreal2sim" / "motion" / "utils"))

from utils.compose_config import compose_configs

class MotionAgent:
    def __init__(self, stage=None, key=None):
        print('[Info] Initializing MotionAgent...')
        self.base_dir = Path.cwd()
        cfg_path = self.base_dir / "config" / "config.yaml"
        cfg = yaml.safe_load(cfg_path.open("r"))
        self.keys = [key] if key is not None else cfg["keys"]
        self.key_cfgs = {key: compose_configs(key, cfg) for key in self.keys}
        self.key_scene_dicts = {}
        for key in self.keys:
            scene_pkl = self.base_dir / f'outputs/{key}/scene/scene.pkl'
            with open(scene_pkl, 'rb') as f:
                scene_dict = pickle.load(f)
            self.key_scene_dicts[key] = scene_dict
        self.key_scene_json_dicts = {}
        for key in self.keys:
            scene_json_path = self.base_dir / f'outputs/{key}/simulation/scene.json'
            with open(scene_json_path, "r") as f:
                scene_json_dict = json.load(f)
            self.key_scene_json_dicts[key] = scene_json_dict
        self.stages = [
            "hand_extraction",
            "demo_motion_process",
            "grasp_point_extraction",
            "grasp_generation",
        ]
        if stage is not None:
            if stage in self.stages:
                start_idx = self.stages.index(stage)
                self.stages = self.stages[start_idx:]
            else:
                print(f"[Warning] Stage '{stage}' not found. It must be one of {self.stages}. Running all stages by default.")
        print('[Info] MotionAgent initialized.')

    def save_scene_dicts(self):
        for key, scene_dict in self.key_scene_dicts.items():
            with open(self.base_dir / f'outputs/{key}/scene/scene.pkl', 'wb') as f:
                pickle.dump(scene_dict, f)
        print('[Info] Scene dictionaries saved.')

    def hand_extraction(self):
        from modules.hand_extraction import hand_extraction
        self.key_scene_dicts = hand_extraction(self.keys, self.key_scene_dicts, self.key_cfgs)
        print("[Info] Hand extraction completed.")

    def demo_motion_process(self):
        from modules.demo_motion_process import demo_motion_process
        self.key_scene_dicts = demo_motion_process(self.keys, self.key_scene_dicts, self.key_cfgs)
        print("[Info] Demo motion process completed.")

    def grasp_point_extraction(self):
        from modules.grasp_point_extraction import grasp_point_extraction
        self.key_scene_dicts = grasp_point_extraction(self.keys, self.key_scene_dicts, self.key_scene_json_dicts, self.key_cfgs)
        print("[Info] Grasp point extraction completed.")

    def grasp_generation(self):
        from modules.grasp_generation import grasp_generation
        grasp_generation(self.keys)
        print("[Info] Grasp generation completed.")

    def run(self):
        if "hand_extraction" in self.stages:
            self.hand_extraction()
        if "demo_motion_process" in self.stages:
            self.demo_motion_process()
        if "grasp_point_extraction" in self.stages:
            self.grasp_point_extraction()
        if "grasp_generation" in self.stages:
            self.grasp_generation()
        self.save_scene_dicts()
        print('[Info] MotionAgent run completed.')
        return self.key_scene_dicts

if __name__ == '__main__':
    args = argparse.ArgumentParser()
    args.add_argument('--stage', type=str, default=None, help='Starting from a certain stage')
    args.add_argument('--key', type=str, default=None, help='Process a single key instead of all keys from config')
    args = args.parse_args()

    agent = MotionAgent(stage=args.stage, key=args.key)
    scene_dicts = agent.run()