success

runners/inference_engine.py

@@ -0,0 +1,132 @@
+import os
+import sys
+from datetime import datetime
+
+import torch
+import pickle
+from tqdm import tqdm
+
+path = os.path.abspath(__file__)
+for i in range(2):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+
+from configs.config import ConfigManager
+from datasets.dataset_factory import DatasetFactory
+from modules.pipeline import Pipeline
+from runners.runner import Runner
+
+
+class InferenceEngine(Runner):
+    RESULTS_DIR_NAME: str = 'results'
+    LOG_DIR_NAME: str = 'log'
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+
+        ''' Pipeline '''
+        self.pipeline_config = ConfigManager.get("settings", "pipeline")
+        self.pipeline = Pipeline(self.pipeline_config).to(self.device)
+
+        ''' Experiment '''
+        self.model_path = ConfigManager.get("settings", "experiment", "model_path")
+        self.load_checkpoint(self.model_path)
+        self.load_experiment("inference")
+        
+        ''' Inference Results '''
+        self.inference_results_config = ConfigManager.get("settings", "results")
+        self.save_data_keys = self.inference_results_config["save_data_keys"]
+        self.save_output_keys = self.inference_results_config["save_output_keys"]
+
+        ''' Test '''
+        self.test_config = ConfigManager.get("settings", "test")
+        self.test_dataset_config_list = self.test_config["dataset_list"]
+        self.test_set_list = []
+        seen_name = set()
+        for test_dataset_config in self.test_dataset_config_list:
+            if test_dataset_config["name"] not in seen_name:
+                seen_name.add(test_dataset_config["name"])
+            else:
+                raise ValueError("Duplicate test dataset name: {}".format(test_dataset_config["name"]))
+            test_set = DatasetFactory.create(test_dataset_config)
+            self.test_set_list.append(test_set)
+        del seen_name
+
+        self.print_info()
+
+    def run(self):
+        print("Inference start...")
+        self.test()
+        print("Inference finished!")
+
+    def test(self):
+        self.pipeline.eval()
+        with torch.no_grad():
+            for dataset_idx, test_set in enumerate(self.test_set_list):
+                test_set_name = self.test_dataset_config_list[dataset_idx]["name"]
+                ratio = self.test_dataset_config_list[dataset_idx]["ratio"]
+
+                test_loader = test_set.get_loader()
+                loop = tqdm(enumerate(test_loader), total=int(len(test_loader)))
+                for i, data in loop:
+                    test_set.process_batch(data, self.device)
+                    output = self.pipeline(data, Pipeline.TEST_MODE)
+                    self.save_output(output, data, test_set_name, i)
+                    loop.set_description(
+                        f'Inference (Test: {test_set_name}, ratio={ratio})')
+                    
+    def save_output(self, output, data, test_set_name, idx):
+        results_dir = os.path.join(str(self.experiment_path), InferenceEngine.RESULTS_DIR_NAME)
+        if not os.path.exists(os.path.join(results_dir,test_set_name)):
+            os.makedirs(os.path.join(results_dir,test_set_name))
+        save_path = os.path.join(results_dir, test_set_name, f"{idx}.pkl")
+        data = {key: value for key, value in data.items() if key in self.save_data_keys}
+        output = {key: value for key, value in output.items() if key in self.save_output_keys}
+        output_converted = {key: value.cpu().numpy() if torch.is_tensor(value) else value for key, value in output.items()}
+        data_converted = {key: value.cpu().numpy() if torch.is_tensor(value) else value for key, value in data.items()}
+        with open(save_path, "wb") as f:
+            pickle.dump({"output":output_converted,"data":data_converted}, f)
+        
+    def load_checkpoint(self, model_path):
+        self.pipeline.load(model_path)
+        print(f"Checkpoint loaded from {model_path}")
+
+    def load_experiment(self, backup_name=None):
+        super().load_experiment(backup_name)
+        
+    def create_experiment(self, backup_name=None):
+        super().create_experiment(backup_name)
+        results_dir = os.path.join(str(self.experiment_path), InferenceEngine.RESULTS_DIR_NAME)
+        os.makedirs(results_dir)
+        
+
+    def print_info(self):
+        def print_dataset(config, dataset):
+            print("\t name: {}".format(config["name"]))
+            print("\t source: {}".format(config["source"]))
+            print("\t data_type: {}".format(config["data_type"]))
+            print("\t total_length: {}".format(len(dataset)))
+            print("\t ratio: {}".format(config["ratio"]))
+            print()
+
+        super().print_info()
+        table_size = 70
+        print(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+')
+        print(self.pipeline)
+        print(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
+        for i, test_dataset_config in enumerate(self.test_dataset_config_list):
+            print(f"test dataset {i}: ")
+            print_dataset(test_dataset_config, self.test_set_list[i])
+
+        print(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="configs/local_inference_config.yaml")
+    args = parser.parse_args()
+    infenrence_engine = InferenceEngine(args.config)
+    infenrence_engine.run()

runners/preprocessor.py

@@ -0,0 +1,71 @@
+import os
+from abc import ABC, abstractmethod
+import shutil
+
+from configs.config import ConfigManager
+from runners.runner import Runner
+
+
+class Preprocessor(Runner, ABC):
+    DATA = "data"
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+
+        self.preprocess_config = ConfigManager.get("settings", "preprocess")
+
+    def load_experiment(self,backup_name=None):
+        super().load_experiment(backup_name)
+        exists_ok = self.experiments_config["keep_exists"]
+        if not exists_ok:
+            data_dir = os.path.join(str(self.experiment_path), Preprocessor.DATA)
+            shutil.rmtree(data_dir, ignore_errors=True)
+            os.makedirs(data_dir)
+            self.create_dataset_list()
+
+    def create_experiment(self,backup_name=None):
+        super().create_experiment(backup_name)
+        data_dir = os.path.join(str(self.experiment_path), Preprocessor.DATA)
+        os.makedirs(data_dir)
+        self.create_dataset_list()
+
+    def create_dataset_list(self):
+        dataset_list = self.preprocess_config["dataset_list"]
+        exists_ok = self.experiments_config["keep_exists"]
+        for dataset in dataset_list:
+            source = dataset["source"]
+            source_dir = os.path.join(str(self.experiment_path), Preprocessor.DATA, source)
+            if not os.path.exists(source_dir):
+                os.makedirs(source_dir,exist_ok=exists_ok)
+            dataset_name = dataset["data_type"]
+            dataset_dir = os.path.join(source_dir, dataset_name)
+            if not os.path.exists(dataset_dir):
+                os.makedirs(dataset_dir,exist_ok=exists_ok)
+
+    @abstractmethod
+    def get_dataloader(self, dataset_config):
+        pass
+
+    @abstractmethod
+    def get_model(self, model_config):
+        pass
+
+    @abstractmethod
+    def prediction(self, model, dataloader):
+        pass
+
+    @abstractmethod
+    def preprocess(self, predicted_data):
+        pass
+
+    @abstractmethod
+    def save_processed_data(self, processed_data, data_config=None):
+        pass
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="../configs/local_gsnet_preprocess_config.yaml")
+    args = parser.parse_args()
+    preproc = Preprocessor(args.config)

runners/preprocessors/grasping/GSNet_preprocessor.py

@@ -0,0 +1,409 @@
+import os
+import re
+import sys
+import numpy as np
+import torch
+import open3d as o3d
+from torch.utils.data import DataLoader
+
+path = os.path.abspath(__file__)
+for i in range(4):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+GSNET_PROJECT_ROOT = os.path.join(PROJECT_ROOT, "baselines/grasping/GSNet")
+sys.path.append(os.path.join(GSNET_PROJECT_ROOT, "pointnet2"))
+sys.path.append(os.path.join(GSNET_PROJECT_ROOT, "utils"))
+sys.path.append(os.path.join(GSNET_PROJECT_ROOT, "models"))
+sys.path.append(os.path.join(GSNET_PROJECT_ROOT, "dataset"))
+
+from utils.omni_util import OmniUtil
+from utils.view_util import ViewUtil
+from runners.preprocessors.grasping.abstract_grasping_preprocessor import GraspingPreprocessor
+from configs.config import ConfigManager
+
+from baselines.grasping.GSNet.models.graspnet import GraspNet
+from baselines.grasping.GSNet.graspnetAPI.graspnetAPI.graspnet_eval import GraspGroup
+from baselines.grasping.GSNet.dataset.graspnet_dataset import minkowski_collate_fn
+from torch.utils.data import Dataset
+
+
+class GSNetInferenceDataset(Dataset):
+    CAMERA_PARAMS_TEMPLATE = "camera_params_{}.json"
+    DISTANCE_TEMPLATE = "distance_to_camera_{}.npy"
+    RGB_TEMPLATE = "rgb_{}.png"
+    MASK_TEMPLATE = "semantic_segmentation_{}.png"
+    MASK_LABELS_TEMPLATE = "semantic_segmentation_labels_{}.json"
+
+    def __init__(
+        self,
+        source="nbv1",
+        data_type="sample",
+        data_dir="/mnt/h/AI/Datasets",
+        scene_pts_num=15000,
+        voxel_size=0.005,
+    ):
+
+        self.data_dir = data_dir
+        self.scene_pts_num = scene_pts_num
+        self.data_path = str(os.path.join(self.data_dir, source, data_type))
+        self.scene_list = os.listdir(self.data_path)
+        self.data_list = self.get_datalist()
+        self.voxel_size = voxel_size
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        frame_path, target = self.data_list[index]
+        frame_data = self.load_frame_data(frame_path=frame_path, object_name=target)
+        return frame_data
+
+    def get_datalist(self):
+        scene_frame_list = []
+        for scene in self.scene_list:
+            scene_path = os.path.join(self.data_path, scene)
+            file_list = os.listdir(scene_path)
+            for file in file_list:
+                if file.startswith("camera_params"):
+                    frame_index = re.findall(r"\d+", file)[0]
+                    frame_path = os.path.join(scene_path, frame_index)
+                    target_list = OmniUtil.get_object_list(frame_path)
+                    for target in target_list:
+                        scene_frame_list.append((frame_path,target))
+                    if len(target_list) == 0:
+                        scene_frame_list.append((frame_path, None))
+            print("Scene: ", scene, " has ", len(scene_frame_list), " frames")
+        return scene_frame_list
+
+    def load_frame_data(self, frame_path, object_name):
+        try:
+            target_list = OmniUtil.get_object_list(path=frame_path, contains_non_obj=True)
+            _, obj_pcl_dict = OmniUtil.get_segmented_points(
+                path=frame_path, target_list=target_list
+            )
+            obj_center = ViewUtil.get_object_center_from_pts_dict(object_name, obj_pcl_dict)
+            croped_pts_dict = ViewUtil.crop_pts_dict(obj_pcl_dict, obj_center, radius=0.2)
+            sampled_scene_pts, sampled_pts_dict = GSNetInferenceDataset.sample_dict_to_target_points(croped_pts_dict)
+            ret_dict  = {
+                "frame_path": frame_path,
+                "point_clouds": sampled_scene_pts.astype(np.float32),
+                "coors": sampled_scene_pts.astype(np.float32) / self.voxel_size,
+                "feats": np.ones_like(sampled_scene_pts).astype(np.float32),
+                "obj_pcl_dict": sampled_pts_dict,
+                "object_name": object_name,
+            }
+        except Exception as e:
+            print("Error in loading frame data: ", e)
+            ret_dict = {
+                "frame_path": frame_path,
+                "point_clouds": np.zeros((self.scene_pts_num, 3)).astype(np.float32),
+                "coors": np.zeros((self.scene_pts_num, 3)).astype(np.float32),
+                "feats": np.ones((self.scene_pts_num, 3)).astype(np.float32),
+                "obj_pcl_dict": {},
+                "object_name": object_name,
+                "error": True
+            }
+        return ret_dict
+    
+    def sample_points(points, target_num_points):
+        num_points = points.shape[0]
+        if num_points == 0:
+            return np.zeros((target_num_points, points.shape[1]))
+        if num_points > target_num_points:
+            indices = np.random.choice(num_points, target_num_points, replace=False)
+        else:
+            indices = np.random.choice(num_points, target_num_points, replace=True)
+        return points[indices]
+
+    def sample_dict_to_target_points(croped_pts_dict, total_points=15000):
+        all_sampled_points = []
+        sampled_pts_dict = {}
+        total_existing_points = sum([pts.shape[0] for pts in croped_pts_dict.values() if pts.shape[0] > 0])
+
+        if total_existing_points > total_points:
+            ratios = {name: len(pts) / total_existing_points for name, pts in croped_pts_dict.items() if pts.shape[0] > 0}
+            target_num_points = {name: int(ratio * total_points) for name, ratio in ratios.items()}
+            remaining_points = total_points - sum(target_num_points.values())
+            for name in target_num_points.keys():
+                if remaining_points > 0:
+                    target_num_points[name] += 1
+                    remaining_points -= 1
+        else:
+            target_num_points = {name: len(pts) for name, pts in croped_pts_dict.items()}
+            remaining_points = total_points - total_existing_points
+            additional_points = np.random.choice([name for name, pts in croped_pts_dict.items() if pts.shape[0] > 0], remaining_points, replace=True)
+            for name in additional_points:
+                target_num_points[name] += 1
+
+        for name, pts in croped_pts_dict.items():
+            if pts.shape[0] == 0:
+                sampled_pts_dict[name] = pts
+                continue
+            sampled_pts = GSNetInferenceDataset.sample_points(pts, target_num_points[name])
+            sampled_pts_dict[name] = sampled_pts
+            all_sampled_points.append(sampled_pts)
+
+        if len(all_sampled_points) > 0:
+            sampled_scene_pts = np.concatenate(all_sampled_points, axis=0)
+        else:
+            sampled_scene_pts = np.zeros((total_points, 3))
+        return sampled_scene_pts, sampled_pts_dict
+
+    @staticmethod
+    def sample_pcl(pcl, n_pts=1024):
+        indices = np.random.choice(pcl.shape[0], n_pts, replace=pcl.shape[0] < n_pts)
+        return pcl[indices, :]
+
+
+class GSNetPreprocessor(GraspingPreprocessor):
+    GRASP_MAX_WIDTH = 0.1
+    GRASPNESS_THRESHOLD = 0.1
+    NUM_VIEW = 300
+    NUM_ANGLE = 12
+    NUM_DEPTH = 4
+    M_POINT = 1024
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+
+    def get_dataloader(self, dataset_config):
+        def my_worker_init_fn(worker_id):
+            np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+        dataset = GSNetInferenceDataset(
+            source=dataset_config["source"],
+            data_type=dataset_config["data_type"],
+            data_dir=dataset_config["data_dir"],
+            scene_pts_num=dataset_config["scene_pts_num"],
+            voxel_size=dataset_config["voxel_size"],
+        )
+        print("Test dataset length: ", len(dataset))
+        dataloader = DataLoader(
+            dataset,
+            batch_size=dataset_config["batch_size"],
+            shuffle=False,
+            num_workers=0,
+            worker_init_fn=my_worker_init_fn,
+            collate_fn=minkowski_collate_fn,
+        )
+        print("Test dataloader length: ", len(dataloader))
+        return dataloader
+
+    def get_model(self, model_config=None):
+        model = GraspNet(seed_feat_dim=model_config["general"]["seed_feat_dim"], is_training=False)
+        model.to("cuda")
+        checkpoint = torch.load(model_config["general"]["checkpoint_path"])
+        model.load_state_dict(checkpoint["model_state_dict"])
+        start_epoch = checkpoint["epoch"]
+        print(
+            "-> loaded checkpoint %s (epoch: %d)" % (model_config["general"]["checkpoint_path"], start_epoch)
+        )
+        model.eval()
+        return model
+
+    def prediction(self, model, dataloader, require_gripper=False, top_k=10):
+        preds = {}
+
+        for idx, batch_data in enumerate(dataloader):
+            try:
+                if "error" in batch_data:
+                    frame_path = batch_data["frame_path"][0]
+                    object_name = batch_data["object_name"][0]
+                    preds[frame_path] = {object_name: None}
+                    print("No graspable points found at frame: ", frame_path)
+                    continue
+                print("Processing batch: ", idx, "/", len(dataloader))
+                for key in batch_data:
+                    if "list" in key:
+                        for i in range(len(batch_data[key])):
+                            for j in range(len(batch_data[key][i])):
+                                batch_data[key][i][j] = batch_data[key][i][j].to("cuda")
+                    elif not isinstance(batch_data[key], (list)):
+                        batch_data[key] = batch_data[key].to("cuda")
+                with torch.no_grad():
+                    
+                    end_points = model(batch_data)
+                    if end_points is None:
+                        frame_path = batch_data["frame_path"][0]
+                        object_name = batch_data["object_name"][0]
+                        preds[frame_path] = {object_name: None}
+                        print("No graspable points found at frame: ", frame_path)
+                        continue
+                    grasp_preds = self.decode_pred(end_points)
+                    
+                    standard_grasp_preds = GSNetPreprocessor.standard_pred_decode(end_points)
+                    standard_preds = standard_grasp_preds[0].detach().cpu().numpy()
+                    if require_gripper:
+                        gg = GraspGroup(standard_preds)
+                        gg = gg.nms()
+                        gg = gg.sort_by_score()
+                        grippers = gg.to_open3d_geometry_list()
+                        gp_pts_list = np.asarray([np.asarray(gripper_mesh.sample_points_uniformly(48).points) for gripper_mesh in grippers], dtype=np.float16)
+                        gp_score_list = gg.scores
+                    
+                    for idx in range(len(batch_data["frame_path"])):
+                        frame_path = batch_data["frame_path"][idx]
+                        object_name = batch_data["object_name"][idx]
+                        if frame_path not in preds:
+                            preds[frame_path] = {object_name: {}}
+                            
+                        preds[frame_path][object_name] = grasp_preds[idx]
+                        preds[frame_path][object_name]["obj_pcl_dict"] = (
+                            batch_data["obj_pcl_dict"][idx]
+                        )
+                        if require_gripper:
+                            preds[frame_path][object_name]["gripper"] = {
+                                "gripper_pose": gp_pts_list.tolist(),
+                                "gripper_score": gp_score_list.tolist()
+                            }
+            except Exception as e:
+                print("Error in inference: ", e)
+                # ----- Debug Trace ----- #
+                print(batch_data["frame_path"])
+                import ipdb; ipdb.set_trace()
+                frame_path = batch_data["frame_path"][idx]
+                object_name = batch_data["object_name"][idx]
+                preds[frame_path] = {object_name: {}}
+                # ------------------------ #
+                  
+
+        results = {}
+        for frame_path in preds:
+            try:
+                predict_results = {}
+                for object_name in preds[frame_path]:
+                    if object_name is None or preds[frame_path][object_name] == None:
+                        continue
+                    grasp_center = preds[frame_path][object_name]["grasp_center"]
+                    grasp_score = preds[frame_path][object_name]["grasp_score"]
+                    obj_pcl_dict = preds[frame_path][object_name]["obj_pcl_dict"]
+                    if require_gripper:
+                        gripper = preds[frame_path][object_name]["gripper"]
+                    grasp_center = grasp_center.unsqueeze(1)
+                    obj_pcl = obj_pcl_dict[object_name]
+                    obj_pcl = torch.tensor(
+                        obj_pcl.astype(np.float32), device=grasp_center.device
+                    )
+                    obj_pcl = obj_pcl.unsqueeze(0)
+                    grasp_obj_table = (grasp_center == obj_pcl).all(axis=-1)
+                    obj_pts_on_grasp = grasp_obj_table.any(axis=1)
+                    obj_graspable_pts = grasp_center[obj_pts_on_grasp].squeeze(1)
+                    
+                    
+                    
+                    obj_graspable_pts_score = grasp_score[obj_pts_on_grasp]
+                    obj_graspable_pts_info = torch.cat(
+                        [obj_graspable_pts, obj_graspable_pts_score], dim=1
+                    )
+                    
+                    if obj_graspable_pts.shape[0] == 0:
+                        obj_graspable_pts_info = torch.zeros((top_k, 4))
+                    ranked_obj_graspable_pts_info = self.sample_graspable_pts(
+                        obj_graspable_pts_info, top_k=top_k
+                    )
+                    predict_results[object_name] = {
+                        "positions": ranked_obj_graspable_pts_info[:, :3]
+                        .cpu()
+                        .numpy()
+                        .tolist(),
+                        "scores": ranked_obj_graspable_pts_info[:, 3]
+                        .cpu()
+                        .numpy()
+                        .tolist(),
+                        
+                    }
+                if require_gripper:
+                    results[frame_path] = {"predicted_results": predict_results, "gripper": gripper}
+                else:
+                    results[frame_path] = {"predicted_results": predict_results}
+
+            except Exception as e:
+                print("Error in postprocessing: ", e)
+                # ----- Debug Trace ----- #
+                print(frame_path)
+                import ipdb; ipdb.set_trace()
+                # ------------------------ #
+                
+        print("Prediction finished")
+        return results
+
+    
+
+    @staticmethod
+    def sample_graspable_pts(graspable_pts, top_k=50):
+        if graspable_pts.shape[0] < top_k:
+            sampled_indices = torch.randint(0, graspable_pts.shape[0], (top_k,))
+            graspable_pts = graspable_pts[sampled_indices]
+        sorted_indices = torch.argsort(graspable_pts[:, 3], descending=True)
+        sampled_indices = graspable_pts[sorted_indices][:top_k]
+        return sampled_indices
+
+    def decode_pred(self, end_points):
+        batch_size = len(end_points["point_clouds"])
+        grasp_preds = []
+        for i in range(batch_size):
+            grasp_center = end_points["xyz_graspable"][i].float()
+            num_pts = end_points["xyz_graspable"][i].shape[0]
+            grasp_score = end_points["grasp_score_pred"][i].float()
+            grasp_score = grasp_score.view(num_pts, -1)
+            grasp_score, _ = torch.max(grasp_score, -1)  # [M_POINT]
+            grasp_score = grasp_score.view(-1, 1)
+            grasp_preds.append(
+                {"grasp_center": grasp_center, "grasp_score": grasp_score}
+            )
+        return grasp_preds
+
+    @staticmethod
+    def standard_pred_decode(end_points):
+        batch_size = len(end_points['point_clouds'])
+        grasp_preds = []
+        for i in range(batch_size):
+            grasp_center = end_points['xyz_graspable'][i].float()
+            num_pts = end_points["xyz_graspable"][i].shape[0]
+            grasp_score = end_points['grasp_score_pred'][i].float()
+            grasp_score = grasp_score.view(num_pts, -1)
+            grasp_score, grasp_score_inds = torch.max(grasp_score, -1)  # [M_POINT]
+            grasp_score = grasp_score.view(-1, 1)
+            grasp_angle = (grasp_score_inds // GSNetPreprocessor.NUM_DEPTH) * np.pi / 12
+            grasp_depth = (grasp_score_inds % GSNetPreprocessor.NUM_DEPTH + 1) * 0.01
+            grasp_depth = grasp_depth.view(-1, 1)
+            grasp_width = 1.2 * end_points['grasp_width_pred'][i] / 10.
+            grasp_width = grasp_width.view(GSNetPreprocessor.M_POINT, GSNetPreprocessor.NUM_ANGLE*GSNetPreprocessor.NUM_DEPTH)
+            grasp_width = torch.gather(grasp_width, 1, grasp_score_inds.view(-1, 1))
+            grasp_width = torch.clamp(grasp_width, min=0., max=GSNetPreprocessor.GRASP_MAX_WIDTH)
+
+            approaching = -end_points['grasp_top_view_xyz'][i].float()
+            grasp_rot = GSNetPreprocessor.batch_viewpoint_params_to_matrix(approaching, grasp_angle)
+            grasp_rot = grasp_rot.view(GSNetPreprocessor.M_POINT, 9)
+
+            # merge preds
+            grasp_height = 0.02 * torch.ones_like(grasp_score)
+            obj_ids = -1 * torch.ones_like(grasp_score)
+            grasp_preds.append(
+                torch.cat([grasp_score, grasp_width, grasp_height, grasp_depth, grasp_rot, grasp_center, obj_ids], axis=-1))
+        return grasp_preds
+    
+    @staticmethod
+    def batch_viewpoint_params_to_matrix(batch_towards, batch_angle):
+        axis_x = batch_towards
+        ones = torch.ones(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device)
+        zeros = torch.zeros(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device)
+        axis_y = torch.stack([-axis_x[:, 1], axis_x[:, 0], zeros], dim=-1)
+        mask_y = (torch.norm(axis_y, dim=-1) == 0)
+        axis_y[mask_y, 1] = 1
+        axis_x = axis_x / torch.norm(axis_x, dim=-1, keepdim=True)
+        axis_y = axis_y / torch.norm(axis_y, dim=-1, keepdim=True)
+        axis_z = torch.cross(axis_x, axis_y)
+        sin = torch.sin(batch_angle)
+        cos = torch.cos(batch_angle)
+        R1 = torch.stack([ones, zeros, zeros, zeros, cos, -sin, zeros, sin, cos], dim=-1)
+        R1 = R1.reshape([-1, 3, 3])
+        R2 = torch.stack([axis_x, axis_y, axis_z], dim=-1)
+        batch_matrix = torch.matmul(R2, R1)
+        return batch_matrix
+
+
+if __name__ == "__main__":
+    gs_preproc = GSNetPreprocessor(config_path="configs/server_gsnet_preprocess_config.yaml")
+    gs_preproc.run()

runners/preprocessors/grasping/abstract_grasping_preprocessor.py

@@ -0,0 +1,65 @@
+import os
+import json
+import numpy as np
+from abc import abstractmethod, ABC
+
+from runners.preprocessor import Preprocessor
+from utils.omni_util import OmniUtil
+
+class GraspingPreprocessor(Preprocessor, ABC):
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        self.load_experiment("GSNet")
+        self.dataset_list_config = self.preprocess_config["dataset_list"]
+        self.model_config = self.preprocess_config["model"]
+
+    def run(self):
+        """
+        - for each dataset
+        --- get its dataloader
+        --- for each batch, do prediction
+        --- preprocess the collected results
+        --- save processed results
+        """
+        for dataset_config in self.dataset_list_config:
+            dataloader = self.get_dataloader(dataset_config)
+            model = self.get_model(self.model_config)
+            predicted_data = self.prediction(model, dataloader)
+            processed_data = self.preprocess(predicted_data)
+            self.save_processed_data(processed_data,dataset_config)
+            
+    def preprocess(self, predicted_data, require_gripper=False):
+        for frame_path in predicted_data:
+            frame_obj_info = predicted_data[frame_path]["predicted_results"]
+            if require_gripper:
+                gripper = predicted_data[frame_path]["gripper"]
+                predicted_data[frame_path]["gripper"] = gripper
+            predicted_data[frame_path]["sum_score"] = {}
+            predicted_data[frame_path]["avg_score"] = {}
+            
+            for obj_name in frame_obj_info:
+                obj_score_sum = np.sum(frame_obj_info[obj_name]["scores"])
+                obj_score_avg = np.mean(frame_obj_info[obj_name]["scores"])
+                predicted_data[frame_path]["sum_score"][obj_name] = obj_score_sum
+                predicted_data[frame_path]["avg_score"][obj_name] = obj_score_avg
+                
+
+        return predicted_data
+
+    def save_processed_data(self, processed_data, data_config=None):
+        data_path = os.path.join(str(self.experiment_path), Preprocessor.DATA, data_config["source"], data_config["data_type"])
+        for frame_path in processed_data:
+            data_item = processed_data[frame_path]
+            scene = os.path.basename(os.path.dirname(frame_path))
+            idx = os.path.basename(frame_path)
+            target_scene_path = os.path.join(str(data_path), scene) 
+            if not os.path.exists(target_scene_path):
+                
+                os.makedirs(target_scene_path)
+            label_save_path = os.path.join(
+                target_scene_path,OmniUtil.SCORE_LABEL_TEMPLATE.format(idx)
+            )
+            with open(label_save_path, "w+") as f:
+                json.dump(data_item, f)
+        print("Processed data saved to: ", data_path)

runners/preprocessors/object_pose/FoundationPose_preprocessor.py

@@ -0,0 +1,185 @@
+import os
+import re
+import sys
+import numpy as np
+import torch
+import trimesh
+from torch.utils.data import DataLoader
+
+
+path = os.path.abspath(__file__)
+for i in range(4):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+
+from utils.omni_util import OmniUtil
+from utils.view_util import ViewUtil
+from runners.preprocessors.object_pose.abstract_object_pose_preprocessor import ObjectPosePreprocessor
+from configs.config import ConfigManager
+
+from torch.utils.data import Dataset
+
+
+class ObjectPoseInferenceDataset(Dataset):
+    CAMERA_PARAMS_TEMPLATE = "camera_params_{}.json"
+    DISTANCE_TEMPLATE = "distance_to_camera_{}.npy"
+    RGB_TEMPLATE = "rgb_{}.png"
+    MASK_TEMPLATE = "semantic_segmentation_{}.png"
+    MASK_LABELS_TEMPLATE = "semantic_segmentation_labels_{}.json"
+
+    def __init__(
+        self,
+        source="nbv1",
+        data_type="sample",
+        data_dir="/mnt/h/AI/Datasets",
+    ):
+
+        self.data_dir = data_dir
+        self.empty_frame = set()
+        self.data_path = str(os.path.join(self.data_dir, source, data_type))
+        self.scene_list = os.listdir(self.data_path)
+        self.data_list = self.get_datalist()
+        
+        self.object_data_list = self.get_object_datalist()
+        self.object_name_list = list(self.object_data_list.keys())
+        self.mesh_dir_path = os.path.join(self.data_dir, source,  "objects")
+        
+        self.meshes = {}
+        self.load_all_meshes()
+        
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        frame_path, target = self.data_list[index]
+        frame_data = self.load_frame_data(frame_path=frame_path, object_name=target)
+        return frame_data
+
+    def load_all_meshes(self):
+        object_name_list = os.listdir(self.mesh_dir_path)
+        for object_name in object_name_list:
+            mesh_path = os.path.join(self.mesh_dir_path, object_name, "Scan", "Simp.obj")
+            mesh = trimesh.load(mesh_path)
+            object_model_scale = [0.001, 0.001, 0.001]
+            mesh.apply_scale(object_model_scale)
+            self.meshes[object_name] = mesh
+    
+    def get_datalist(self):
+        for scene in self.scene_list:
+            scene_path = os.path.join(self.data_path, scene)
+            file_list = os.listdir(scene_path)
+            scene_frame_list = []
+            for file in file_list:
+                if file.startswith("camera_params"):
+                    frame_index = re.findall(r"\d+", file)[0]
+                    frame_path = os.path.join(scene_path, frame_index)
+                    target_list = OmniUtil.get_object_list(frame_path)
+                    for target in target_list:
+                        scene_frame_list.append((frame_path,target))
+                    if len(target_list) == 0:
+                        self.empty_frame.add(frame_path)
+
+        return scene_frame_list
+    
+    def get_object_datalist(self):
+        object_datalist = {}
+        for data_item in self.data_list:
+            frame_path, target = data_item
+            if target not in object_datalist:
+                object_datalist[target] = []
+            object_datalist[target].append(frame_path)
+        return object_datalist
+    
+    def get_object_data_batch(self, object_name):
+        object_data_list = self.object_data_list[object_name]
+        batch_data = {"frame_path_list":[],
+                      "rgb_batch":[],
+                      "depth_batch":[],
+                      "seg_batch":[],
+                      "gt_pose_batch":[],
+                      "K":None,
+                      "mesh":None}
+        for frame_path in object_data_list:
+            frame_data = self.load_frame_data(frame_path, object_name)
+            batch_data["frame_path_list"].append(frame_path)
+            batch_data["rgb_batch"].append(frame_data["rgb"])
+            batch_data["depth_batch"].append(frame_data["depth"])
+            batch_data["seg_batch"].append(frame_data["seg"])
+            batch_data["gt_pose_batch"].append(frame_data["gt_pose"])
+            batch_data["K"] = frame_data["K"]
+            batch_data["mesh"] = frame_data["mesh"]
+        
+        batch_data["rgb_batch"] = np.asarray(batch_data["rgb_batch"],dtype=np.uint8)
+        batch_data["depth_batch"] = np.asarray(batch_data["depth_batch"])
+        batch_data["seg_batch"] = np.asarray(batch_data["seg_batch"])
+        batch_data["gt_pose_batch"] = np.asarray(batch_data["gt_pose_batch"])
+        return batch_data
+    
+    def load_frame_data(self, frame_path, object_name):
+        rgb = OmniUtil.get_rgb(frame_path)
+        depth = OmniUtil.get_depth(frame_path)
+        seg = OmniUtil.get_single_seg(frame_path, object_name)
+        K = OmniUtil.get_intrinsic_matrix(frame_path)
+        gt_obj_pose = OmniUtil.get_o2c_pose(frame_path, object_name)
+        ret_dict  = {
+            "frame_path": frame_path,
+            "rgb": rgb.astype(np.float32),
+            "depth": depth.astype(np.float32),
+            "seg": seg,
+            "K": K.astype(np.float32),
+            "object_name": object_name,
+            "mesh": self.meshes[object_name],
+            "gt_pose": gt_obj_pose.astype(np.float32)
+        }
+        return ret_dict
+    
+class FoundationPosePreprocessor(ObjectPosePreprocessor):
+    
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        
+    def run(self):
+        for dataset_config in self.dataset_list_config:
+            dataset = ObjectPoseInferenceDataset( 
+                source=dataset_config["source"],
+                data_type=dataset_config["data_type"],
+                data_dir=dataset_config["data_dir"],
+                )
+            result = self.prediction(dataset)
+            self.save_processed_data(result, dataset_config)
+            
+    def prediction(self, dataset):
+        final_result = {}
+        cnt = 0
+        for object_name in dataset.object_name_list:
+            cnt	+= 1
+            print(f"Processing object: {object_name} ({cnt}/{len(dataset.object_name_list)})")
+            object_data_batch = dataset.get_object_data_batch(object_name)
+            print(f"batch size of object {object_name}: {len(object_data_batch['frame_path_list'])}")
+            pose_batch, result_batch = ViewUtil.get_object_pose_batch(
+                object_data_batch["K"],
+                object_data_batch["mesh"],
+                object_data_batch["rgb_batch"],
+                object_data_batch["depth_batch"],
+                object_data_batch["seg_batch"],
+                object_data_batch["gt_pose_batch"],
+                self.web_server_config["port"]
+            )
+            for frame_path, pred_pose,gt_pose,result in zip(object_data_batch["frame_path_list"], pose_batch,object_data_batch["gt_pose_batch"],result_batch):
+                if frame_path not in final_result:
+                    final_result[frame_path]={}
+                final_result[frame_path][object_name] = {"gt_pose":gt_pose.tolist(),"pred_pose":pred_pose.tolist(),"eval_result":result}
+        for frame_path in dataset.empty_frame:
+            final_result[frame_path] = {}
+        return final_result
+            
+if __name__ == "__main__":
+    config_path = os.path.join(PROJECT_ROOT, "configs/server_object_preprocess_config.yaml")
+    preprocessor = FoundationPosePreprocessor(config_path)
+    preprocessor.run()
+    
+    
+    
+    
+    

runners/preprocessors/object_pose/abstract_object_pose_preprocessor.py

@@ -0,0 +1,51 @@
+import os
+import json
+import numpy as np
+from abc import abstractmethod, ABC
+
+from runners.preprocessor import Preprocessor
+from utils.omni_util import OmniUtil
+
+class ObjectPosePreprocessor(Preprocessor, ABC):
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        self.load_experiment("GSNet")
+        self.dataset_list_config = self.preprocess_config["dataset_list"]
+        self.web_server_config = self.preprocess_config["web_server"]
+
+    
+    def run(self):
+        pass
+    
+   
+    def get_model(self, model_config):
+        pass
+
+    def get_dataloader(self, dataset_config):
+        pass
+    
+    def preprocess(self, predicted_data):
+        pass
+
+    def prediction(self, model, dataloader):
+        pass
+
+    def save_processed_data(self, processed_data, data_config=None):
+        data_path = os.path.join(str(self.experiment_path), Preprocessor.DATA, data_config["source"], data_config["data_type"])
+        # ----- Debug Trace ----- #
+        import ipdb; ipdb.set_trace()
+        # ------------------------ #
+        for frame_path in processed_data:
+            data_item = processed_data[frame_path]
+            scene = os.path.basename(os.path.dirname(frame_path))
+            idx = os.path.basename(frame_path)
+            target_scene_path = os.path.join(str(data_path), scene) 
+            if not os.path.exists(target_scene_path):
+                os.makedirs(target_scene_path)
+            label_save_path = os.path.join(
+                target_scene_path,OmniUtil.SCORE_LABEL_TEMPLATE.format(idx)
+            )
+            with open(label_save_path, "w+") as f:
+                json.dump(data_item, f)
+        print("Processed data saved to: ", data_path)

runners/preprocessors/rgb_feat/abstract_rgb_feat_preprocessor.py

@@ -0,0 +1,47 @@
+import os
+import json
+import numpy as np
+from abc import abstractmethod, ABC
+
+from runners.preprocessor import Preprocessor
+from utils.omni_util import OmniUtil
+
+class RGBFeatPreprocessor(Preprocessor, ABC):
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        self.load_experiment("RGBFeat")
+        self.dataset_list_config = self.preprocess_config["dataset_list"]
+        self.model_config = self.preprocess_config["model"]
+
+    def run(self):
+        """
+        - for each dataset
+        --- get its dataloader
+        --- for each batch, do prediction
+        --- preprocess the collected results
+        --- save processed results
+        """
+        for dataset_config in self.dataset_list_config:
+            dataloader = self.get_dataloader(dataset_config)
+            model = self.get_model(self.model_config)
+            predicted_data = self.prediction(model, dataloader)
+            self.save_processed_data(predicted_data,dataset_config)
+
+    def preprocess(self, predicted_data):
+        pass
+    
+    def save_processed_data(self, processed_data, data_config=None):
+        data_path = os.path.join(str(self.experiment_path), Preprocessor.DATA, data_config["source"], data_config["data_type"])
+        for frame_path in processed_data:
+            rgb_feat = processed_data[frame_path]
+            scene = os.path.basename(os.path.dirname(frame_path))
+            idx = os.path.basename(frame_path)
+            target_scene_path = os.path.join(str(data_path), scene) 
+            if not os.path.exists(target_scene_path):
+                os.makedirs(target_scene_path)
+            rgb_feat_save_path = os.path.join(
+                target_scene_path,OmniUtil.RGB_FEAT_TEMPLATE.format(idx))
+            np.save(rgb_feat_save_path, rgb_feat)
+        print("Processed data saved to: ", data_path)
+

runners/preprocessors/rgb_feat/dinov2_preprocessor.py

@@ -0,0 +1,128 @@
+import os
+import sys
+
+path = os.path.abspath(__file__)
+for i in range(4):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+
+import re
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from utils.omni_util import OmniUtil
+from runners.preprocessors.rgb_feat.abstract_rgb_feat_preprocessor import RGBFeatPreprocessor
+from modules.rgb_encoder.dinov2_encoder import Dinov2Encoder
+from PIL import Image
+
+from torch.utils.data import Dataset
+
+
+class Dinov2InferenceDataset(Dataset):
+    RGB_TEMPLATE = "rgb_{}.png"
+
+    def __init__(
+        self,
+        source="nbv1",
+        data_type="sample",
+        data_dir="/mnt/h/AI/Datasets",
+        image_size = 480
+    ):
+
+        self.data_dir = data_dir
+        self.data_path = str(os.path.join(self.data_dir, source, data_type))
+        self.scene_list = os.listdir(self.data_path)
+        self.data_list = self.get_datalist()
+        self.transform = transforms.Compose([           
+                        transforms.Resize(image_size),
+                        transforms.CenterCrop(int(image_size//14)*14),              
+                        transforms.ToTensor(),                    
+                        transforms.Normalize(mean=0.5, std=0.2)
+                        ])
+
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        frame_path = self.data_list[index]
+        frame_data = self.load_frame_data(frame_path=frame_path)
+        return frame_data
+
+    def get_datalist(self):
+        for scene in self.scene_list:
+            scene_path = os.path.join(self.data_path, scene)
+            file_list = os.listdir(scene_path)
+            scene_frame_list = []
+            for file in file_list:
+                if file.startswith("camera_params"):
+                    frame_index = re.findall(r"\d+", file)[0]
+                    frame_path = os.path.join(scene_path, frame_index)
+                    scene_frame_list.append(frame_path)
+
+        return scene_frame_list
+
+    def load_frame_data(self, frame_path):
+        rgb = OmniUtil.get_rgb(frame_path)
+        rgb = Image.fromarray(rgb)
+        rgb = self.transform(rgb)
+        ret_dict = {"rgb": rgb, "frame_path": frame_path}
+        return ret_dict
+
+
+class Dinov2Preprocessor(RGBFeatPreprocessor):
+    MODULE_NAME: str = "dinov2"
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+
+    def get_dataloader(self, dataset_config):
+
+        dataset = Dinov2InferenceDataset(
+            source=dataset_config["source"],
+            data_type=dataset_config["data_type"],
+            data_dir=dataset_config["data_dir"],
+            image_size = dataset_config["image_size"]
+        )
+        print("Test dataset length: ", len(dataset))
+        dataloader = DataLoader(
+            dataset,
+            batch_size=dataset_config["batch_size"],
+            shuffle=False,
+            num_workers=0,
+        )
+        print("Test dataloader length: ", len(dataloader))
+        return dataloader
+
+    def get_model(self, model_config=None):
+        model = Dinov2Encoder(model_config["general"]["model_name"])
+        model.to("cuda")
+        return model
+
+    def prediction(self, model, dataloader):
+        results = {}
+        total = len(dataloader)
+        for idx, batch_data in enumerate(dataloader):
+            rgb = batch_data["rgb"].to("cuda")
+            with torch.no_grad():
+                rgb_feat = model.encode_rgb(rgb)
+            frame_paths = batch_data["frame_path"]
+            for i, frame_path in enumerate(frame_paths):
+                results[frame_path] = rgb_feat[i].cpu().numpy()
+            print(f"Processed {idx}/{total} batches")
+                
+        return results
+
+    def visualize_feature(self, rgb_feat, model_name, save_path=None):
+        model = Dinov2Encoder(model_name)
+        model.visualize_features(rgb_feat,save_path)
+
+
+if __name__ == "__main__":
+    rgb_preproc = Dinov2Preprocessor(config_path="configs/server_rgb_feat_preprocess_config.yaml")
+    #ßrgb_preproc.run()
+    rgb_feat = np.load("experiments/rgb_feat_preprocessor_test/data/nbv1/sample/scene_0/rgb_feat_0405.npy")
+    
+    rgb_preproc.visualize_feature(rgb_feat, "dinov2_vits14", './visualize.png')

runners/runner.py

@@ -0,0 +1,60 @@
+import os
+import sys
+import time
+
+from abc import abstractmethod, ABC
+import numpy as np
+import torch
+
+from configs.config import ConfigManager
+
+class Runner(ABC):
+
+    @abstractmethod
+    def __init__(self, config_path):
+        ConfigManager.load_config_with(config_path)
+        ConfigManager.print_config()
+        seed = ConfigManager.get("settings", "general", "seed")
+        self.device = ConfigManager.get("settings", "general", "device")
+        self.cuda_visible_devices = ConfigManager.get("settings","general","cuda_visible_devices")
+        os.environ["CUDA_VISIBLE_DEVICES"] = self.cuda_visible_devices
+        self.experiments_config = ConfigManager.get("settings", "experiment")
+        self.experiment_path = os.path.join(self.experiments_config["root_dir"], self.experiments_config["name"])
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        lt = time.localtime()
+        self.file_name = f"{lt.tm_year}_{lt.tm_mon}_{lt.tm_mday}_{lt.tm_hour}h{lt.tm_min}m{lt.tm_sec}s"
+
+    @abstractmethod
+    def run(self):
+        pass
+
+    @abstractmethod
+    def load_experiment(self, backup_name=None):
+        if not os.path.exists(self.experiment_path):
+            print(f"experiments environment {self.experiments_config['name']} does not exists.")
+            self.create_experiment(backup_name)
+        else:
+            print(f"experiments environment {self.experiments_config['name']}")
+            backup_config_dir = os.path.join(str(self.experiment_path), "configs")
+            if not os.path.exists(backup_config_dir):
+                os.makedirs(backup_config_dir)
+            ConfigManager.backup_config_to(backup_config_dir, self.file_name, backup_name)
+
+    @abstractmethod
+    def create_experiment(self, backup_name=None):
+        print("creating experiment: " + self.experiments_config["name"])
+        os.makedirs(self.experiment_path)
+        backup_config_dir = os.path.join(str(self.experiment_path), "configs")
+        os.makedirs(backup_config_dir)
+        ConfigManager.backup_config_to(backup_config_dir, self.file_name, backup_name)
+        log_dir = os.path.join(str(self.experiment_path), "log")
+        os.makedirs(log_dir)
+        cache_dir = os.path.join(str(self.experiment_path), "cache")
+        os.makedirs(cache_dir)
+
+    def print_info(self):
+        table_size = 80
+        print("+" + "-" * table_size + "+")
+        print(f"| Experiment <{self.experiments_config['name']}>")
+        print("+" + "-" * table_size + "+")

runners/tensorboard_runner.py

@@ -0,0 +1,37 @@
+import os
+import subprocess
+import sys
+
+def find_free_port(start_port):
+    import socket
+    port = start_port
+    while True:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            if s.connect_ex(('localhost', port)) != 0:
+                return port
+            port += 1
+def run(exp_name, exp_root="experiments",port=None):
+    port = 6007 if port is None else port
+    max_attempts = 10
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    tensorboard_root = os.path.join(project_root, exp_root, exp_name, "tensorboard")
+    
+    for attempt in range(max_attempts):
+        try:
+            print(f"Trying to launch TensorBoard on port {port}...")
+            subprocess.check_call([
+                sys.executable, "-m", "tensorboard.main",
+                f"--logdir={tensorboard_root}",
+                f"--port={port}"
+            ])
+            break
+        except subprocess.CalledProcessError as e:
+            print(f"Port {port} is in use, trying next port...")
+            port = find_free_port(port + 1)
+    else:
+        print("Failed to launch TensorBoard after multiple attempts.")
+
+if __name__ == "__main__":
+    exp_root = "experiments"
+    exp_name = "sample_train_100_item_overfit_foreground_0"
+    run(exp_name,exp_root,port=6009)

runners/tester.py

@@ -0,0 +1,33 @@
+import os
+
+from configs.config import ConfigManager
+from runners.runner import Runner
+
+
+class Tester(Runner):
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        self.pipeline_config = ConfigManager.get("settings", "pipeline")
+        self.current_epoch = 0
+
+    def run(self):
+        pass
+
+    def load_experiment(self):
+        super().load_experiment()
+
+    def create_experiment(self):
+        super().create_experiment()
+        experiment_path = os.path.join(self.experiments_config["root_dir"], self.experiments_config["name"])
+        result_dir = os.path.join(str(experiment_path), "results")
+        os.makedirs(result_dir)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="configs/local_train_config.yaml")
+    args = parser.parse_args()
+    tester = Tester(args.config)
+    tester.run()

runners/trainer.py

@@ -0,0 +1,257 @@
+import os
+import sys
+from datetime import datetime
+
+import torch
+from tqdm import tqdm
+from torch.utils.tensorboard import SummaryWriter
+
+path = os.path.abspath(__file__)
+for i in range(2):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+
+from configs.config import ConfigManager
+from datasets.dataset_factory import DatasetFactory
+from optimizers.optimizer_factory import OptimizerFactory
+from evaluations.eval_function_factory import EvalFunctionFactory
+from losses.loss_function_factory import LossFunctionFactory
+from modules.pipeline import Pipeline
+from runners.runner import Runner
+from utils.file_util import FileUtil
+from utils.tensorboard_util import TensorboardWriter
+from annotations.external_module import EXTERNAL_FREEZE_MODULES
+
+
+class Trainer(Runner):
+    CHECKPOINT_DIR_NAME: str = 'checkpoints'
+    TENSORBOARD_DIR_NAME: str = 'tensorboard'
+    LOG_DIR_NAME: str = 'log'
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        tensorboard_path = os.path.join(self.experiment_path, Trainer.TENSORBOARD_DIR_NAME)
+
+        ''' Pipeline '''
+        self.pipeline_config = ConfigManager.get("settings", "pipeline")
+        self.parallel = ConfigManager.get("settings","general","parallel")
+        self.pipeline = Pipeline(self.pipeline_config)
+        if self.parallel and self.device == "cuda":
+            self.pipeline = torch.nn.DataParallel(self.pipeline)
+        self.pipeline = self.pipeline.to(self.device)
+            
+        ''' Experiment '''
+        self.current_epoch = 0
+        self.max_epochs = self.experiments_config["max_epochs"]
+        self.test_first = self.experiments_config["test_first"]
+        self.load_experiment("train")
+
+        ''' Train '''
+        self.train_config = ConfigManager.get("settings", "train")
+        self.train_dataset_config = self.train_config["dataset"]
+        self.train_set = DatasetFactory.create(self.train_dataset_config)
+        self.optimizer = OptimizerFactory.create(self.train_config["optimizer"], self.pipeline.parameters())
+        self.train_writer = SummaryWriter(
+            log_dir=os.path.join(tensorboard_path, f"[train]{self.train_dataset_config['name']}"))
+
+        ''' Test '''
+        self.test_config = ConfigManager.get("settings", "test")
+        self.test_dataset_config_list = self.test_config["dataset_list"]
+        self.test_set_list = []
+        self.test_writer_list = []
+        seen_name = set()
+        for test_dataset_config in self.test_dataset_config_list:
+            if test_dataset_config["name"] not in seen_name:
+                seen_name.add(test_dataset_config["name"])
+            else:
+                raise ValueError("Duplicate test dataset name: {}".format(test_dataset_config["name"]))
+            test_set = DatasetFactory.create(test_dataset_config)
+            test_writer = SummaryWriter(
+                log_dir=os.path.join(tensorboard_path, f"[test]{test_dataset_config['name']}"))
+            self.test_set_list.append(test_set)
+            self.test_writer_list.append(test_writer)
+        del seen_name
+
+        self.print_info()
+
+    def run(self):
+        save_interval = self.experiments_config["save_checkpoint_interval"]
+        if self.current_epoch != 0:
+            print("Continue training from epoch {}.".format(self.current_epoch))
+        else:
+            print("Start training from initial model.")
+        if self.test_first:
+            print("Do test first.")
+            self.test()
+        while self.current_epoch < self.max_epochs:
+            self.current_epoch += 1
+            self.train()
+            self.test()
+            if self.current_epoch % save_interval == 0:
+                self.save_checkpoint()
+            self.save_checkpoint(is_last=True)
+
+    def train(self):
+        self.pipeline.train()
+        train_set_name = self.train_dataset_config["name"]
+        ratio = self.train_dataset_config["ratio"]
+        train_loader = self.train_set.get_loader(device="cuda", shuffle=True)
+
+        loop = tqdm(enumerate(train_loader), total=len(train_loader))
+        loader_length = len(train_loader)
+        for i, data in loop:
+            self.train_set.process_batch(data, self.device)
+            loss_dict = self.train_step(data)
+            loop.set_description(
+                f'Epoch [{self.current_epoch}/{self.max_epochs}] (Train: {train_set_name}, ratio={ratio})')
+            loop.set_postfix(loss=loss_dict)
+            curr_iters = (self.current_epoch - 1) * loader_length + i
+            TensorboardWriter.write_tensorboard(self.train_writer, "iter", loss_dict, curr_iters)
+
+    def train_step(self, data):
+        self.optimizer.zero_grad()
+        output = self.pipeline(data, Pipeline.TRAIN_MODE)
+        total_loss, loss_dict = self.loss_fn(output, data)
+        total_loss.backward()
+        self.optimizer.step()
+        for k, v in loss_dict.items():
+            loss_dict[k] = round(v, 5)
+        return loss_dict
+
+    def loss_fn(self, output, data):
+        loss_config = self.train_config["losses"]
+        loss_dict = {}
+        total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.device)
+        for key in loss_config:
+            weight = loss_config[key]
+            target_loss_fn = LossFunctionFactory.create(key)
+            loss = target_loss_fn(output, data)
+            loss_dict[key] = loss.item()
+            total_loss += weight * loss
+
+        loss_dict['total_loss'] = total_loss.item()
+        return total_loss, loss_dict
+
+    def test(self):
+        self.pipeline.eval()
+        with torch.no_grad():
+            for dataset_idx, test_set in enumerate(self.test_set_list):
+                eval_list = self.test_dataset_config_list[dataset_idx]["eval_list"]
+                test_set_name = self.test_dataset_config_list[dataset_idx]["name"]
+                ratio = self.test_dataset_config_list[dataset_idx]["ratio"]
+                writer = self.test_writer_list[dataset_idx]
+                output_list = []
+                data_list = []
+                test_loader = test_set.get_loader("cpu")
+                loop = tqdm(enumerate(test_loader), total=int(len(test_loader)))
+                for i, data in loop:
+                    test_set.process_batch(data, self.device)
+                    output = self.pipeline(data, Pipeline.TEST_MODE)
+                    output_list.append(output)
+                    data_list.append(data)
+                    loop.set_description(
+                        f'Epoch [{self.current_epoch}/{self.max_epochs}] (Test: {test_set_name}, ratio={ratio})')
+                result_dict = self.eval_fn(output_list, data_list, eval_list)
+                TensorboardWriter.write_tensorboard(writer, "epoch", result_dict, self.current_epoch - 1)
+
+    @staticmethod
+    def eval_fn(output_list, data_list, eval_list):
+        target_eval_fn = EvalFunctionFactory.create(eval_list)
+        result_dict = target_eval_fn(output_list, data_list)
+        return result_dict
+
+    def get_checkpoint_path(self, is_last=False):
+        return os.path.join(self.experiment_path, Trainer.CHECKPOINT_DIR_NAME,
+                            "Epoch_{}.pth".format(
+                                self.current_epoch if self.current_epoch != -1 and not is_last else "last"))
+
+    def load_checkpoint(self, is_last=False):
+        self.load(self.get_checkpoint_path(is_last))
+        print(f"Loaded checkpoint from {self.get_checkpoint_path(is_last)}")
+        if is_last:
+            checkpoint_root = os.path.join(self.experiment_path, Trainer.CHECKPOINT_DIR_NAME)
+            meta_path = os.path.join(checkpoint_root, "meta.json")
+            if not os.path.exists(meta_path):
+                raise FileNotFoundError(
+                    "No checkpoint meta.json file in the experiment {}".format(self.experiments_config["name"]))
+            meta = FileUtil.load_json("meta.json", checkpoint_root)
+            self.current_epoch = meta["last_epoch"]
+
+    def save_checkpoint(self, is_last=False):
+        self.save(self.get_checkpoint_path(is_last))
+        if not is_last:
+            print(f"Checkpoint at epoch {self.current_epoch} saved to {self.get_checkpoint_path(is_last)}")
+        else:
+            meta = {
+                "last_epoch": self.current_epoch,
+                "time": str(datetime.now())
+            }
+            checkpoint_root = os.path.join(self.experiment_path, Trainer.CHECKPOINT_DIR_NAME)
+            FileUtil.save_json(meta, "meta.json", checkpoint_root)
+
+    def load_experiment(self, backup_name=None):
+        super().load_experiment(backup_name)
+        if self.experiments_config["use_checkpoint"]:
+            self.current_epoch = self.experiments_config["epoch"]
+            self.load_checkpoint(is_last=(self.current_epoch == -1))
+
+    def create_experiment(self, backup_name=None):
+        super().create_experiment(backup_name)
+        ckpt_dir = os.path.join(str(self.experiment_path), Trainer.CHECKPOINT_DIR_NAME)
+        os.makedirs(ckpt_dir)
+        tensorboard_dir = os.path.join(str(self.experiment_path), Trainer.TENSORBOARD_DIR_NAME)
+        os.makedirs(tensorboard_dir)
+        
+    def load(self, path):
+        state_dict = torch.load(path)
+        if self.parallel:
+            self.pipeline.module.load_state_dict(state_dict)
+        else:
+            self.pipeline.load_state_dict(state_dict)
+
+    def save(self, path):
+        if self.parallel:
+            state_dict = self.pipeline.module.state_dict()
+        else:
+            state_dict = self.pipeline.state_dict()
+
+        for name, module in self.pipeline.named_modules():
+            if module.__class__ in EXTERNAL_FREEZE_MODULES:
+                if name in state_dict:
+                    del state_dict[name]
+
+        torch.save(state_dict, path)
+
+
+    def print_info(self):
+        def print_dataset(config, dataset):
+            print("\t name: {}".format(config["name"]))
+            print("\t source: {}".format(config["source"]))
+            print("\t data_type: {}".format(config["data_type"]))
+            print("\t total_length: {}".format(len(dataset)))
+            print("\t ratio: {}".format(config["ratio"]))
+            print()
+
+        super().print_info()
+        table_size = 70
+        print(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+')
+        print(self.pipeline)
+        print(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
+        print("train dataset: ")
+        print_dataset(self.train_dataset_config, self.train_set)
+        for i, test_dataset_config in enumerate(self.test_dataset_config_list):
+            print(f"test dataset {i}: ")
+            print_dataset(test_dataset_config, self.test_set_list[i])
+
+        print(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="configs/server_train_config.yaml")
+    args = parser.parse_args()
+    trainer = Trainer(args.config)
+    trainer.run()

runners/view_generator.py

@@ -0,0 +1,190 @@
+import os
+import pickle
+import pybullet as p
+import pybullet_data
+import numpy as np
+import math
+from flask import Flask, request, jsonify
+
+import sys
+
+path = os.path.abspath(__file__)
+for i in range(2):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+
+from runners.runner import Runner
+from configs.config import ConfigManager
+
+class ViewGenerator(Runner):
+    def __init__(self, config_path, camera_params) -> None:
+        super().__init__(config_path)
+        self.data_dir = ConfigManager.get("settings", "dataset", "data_dir")
+        self.port = ConfigManager.get("settings", "web_api", "port")
+        self.camera_params = camera_params
+        self.object_model_scale = [0.001, 0.001, 0.001]
+        self.segmentation_labels = {}
+        self.app = Flask(__name__)
+        self._init_routes()
+        
+    def create_experiment(self, backup_name=None):
+        return super().create_experiment(backup_name)
+    
+    def load_experiment(self, backup_name=None):
+        return super().load_experiment(backup_name)
+
+    def _init_routes(self):
+        @self.app.route("/get_images", methods=["POST"])
+        def get_images_api():
+            data = request.get_json()
+            camera_pose = data["camera_pose"]
+            scene = data["scene"]
+            data_type = data["data_type"]
+            source = data["source"]
+            scene_path = os.path.join(self.data_dir, source, data_type, scene)
+            model_dir = os.path.join(self.data_dir, source, "objects")
+            self.load_scene(scene_path,model_dir)
+            result = self.generate_images(camera_pose)
+            result = {
+                "rgb": result["rgb"].tolist(),
+                "depth": result["depth"].tolist(),
+                "segmentation": result["segmentation"].tolist(),
+                "segmentation_labels": result["segmentation_labels"],
+                "camera_params": result["camera_params"],
+            }
+
+            return jsonify(result)
+
+    def load_scene(self, scene_path, model_dir):
+        scene_path = os.path.join(scene_path, "scene.pickle")
+        self.scene = pickle.load(open(scene_path, "rb"))
+        self._initialize_pybullet_scene(model_dir)
+
+    def _initialize_pybullet_scene(self,model_dir):
+        if p.isConnected():
+            p.resetSimulation()
+        else:
+            p.connect(p.DIRECT)
+        p.setAdditionalSearchPath(pybullet_data.getDataPath())
+        p.setGravity(0, 0, 0)
+        p.loadURDF("plane100.urdf")
+        for obj_name in self.scene.keys():
+            orientation = self.scene[obj_name]["rotation"]
+            position = self.scene[obj_name]["position"]
+            class_name = obj_name[:-4]
+            obj_path = os.path.join(model_dir,class_name, obj_name, "Scan", "Simp.obj")
+            self._load_obj_to_pybullet(
+                obj_file_path=obj_path,
+                position=position,
+                orientation=orientation,
+                scale=self.object_model_scale,
+            )
+
+    def _load_obj_to_pybullet(self, obj_file_path, position, orientation, scale):
+        visual_ind = p.createVisualShape(
+            shapeType=p.GEOM_MESH,
+            fileName=obj_file_path,
+            rgbaColor=[1, 1, 1, 1],
+            specularColor=[0.4, 0.4, 0],
+            visualFramePosition=[0, 0, 0],
+            meshScale=scale,
+        )
+        p.createMultiBody(
+            baseMass=1,
+            baseVisualShapeIndex=visual_ind,
+            basePosition=position,
+            baseOrientation=orientation,
+            useMaximalCoordinates=True,
+        )
+
+    def _render_image(self, camera_pose):
+        width = self.camera_params["width"]
+        height = self.camera_params["height"]
+        fov = self.camera_params["fov"]
+        aspect = width / height
+        near = self.camera_params["near"]
+        far = self.camera_params["far"]
+
+        T = np.mat([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 1], [0, 0, 0, 1]])
+        look_at_T = camera_pose.dot(T)
+        view_matrix = p.computeViewMatrix(
+            [camera_pose[0, 3], camera_pose[1, 3], camera_pose[2, 3]],
+            [look_at_T[0, 3], look_at_T[1, 3], look_at_T[2, 3]],
+            [-camera_pose[0, 1], -camera_pose[1, 1], -camera_pose[2, 1]],
+        )
+        projection_matrix = p.computeProjectionMatrixFOV(fov, aspect, near, far)
+        images = p.getCameraImage(
+            width,
+            height,
+            view_matrix,
+            projection_matrix,
+            renderer=p.ER_BULLET_HARDWARE_OPENGL,
+        )
+        rgb = images[2]
+        depth = images[3]
+        seg = images[4]
+        rgb = np.reshape(rgb, (height, width, 4))
+        depth = np.reshape(depth, (height, width))
+        seg = np.reshape(seg, (height, width))
+        rgb_image = rgb[..., :3]
+
+        depth_image = far * near / (far - (far - near) * depth)
+        depth_image = np.asanyarray(depth_image).astype(np.float32) * 1000.0
+        depth_image = depth_image.astype(np.uint16)
+
+        id = 0
+        for object_name in self.scene.keys():
+            self.segmentation_labels[str(id + 1)] = object_name
+            id += 1
+
+        return {
+            "rgb": rgb_image,
+            "depth": depth_image,
+            "segmentation": seg,
+            "segmentation_labels": self.segmentation_labels,
+            "camera_params": self.camera_params,
+        }
+
+    def generate_images(self, camera_pose):
+        results = self._render_image(np.asarray(camera_pose))
+        p.stepSimulation()
+        return results
+
+    def run(self):
+        self.app.run(host="0.0.0.0", port=self.port)
+
+ISAAC_SIM_CAM_H_APERTURE = 20.955
+ISAAC_SIM_CAM_V_APERTURE = 15.2908
+ISAAC_SIM_FOCAL_LENGTH = 39
+ISAAC_SIM_CAM_D_APERTURE = math.sqrt(ISAAC_SIM_CAM_H_APERTURE**2 + ISAAC_SIM_CAM_V_APERTURE**2)
+
+CAM_WIDTH = 640
+CAM_HEIGHT = 480
+CAM_FOV = 2 * math.atan(ISAAC_SIM_CAM_D_APERTURE / (2 * ISAAC_SIM_FOCAL_LENGTH)) / math.pi * 180
+CAM_NEAR = 0.001
+CAM_FAR = 10
+CAM_CX = CAM_WIDTH / 2
+CAM_CY = CAM_HEIGHT / 2
+CAM_FX = 1 / math.tan(CAM_FOV * math.pi / 180.0 / 2.0) * CAM_WIDTH / 2
+CAM_FY = 1 / (CAM_HEIGHT / CAM_WIDTH * math.tan(CAM_FOV * math.pi / 180.0 / 2.0)) * CAM_HEIGHT / 2
+
+CAMERA_PARAMS = {
+    "width": CAM_WIDTH,
+    "height": CAM_HEIGHT,
+    "fov": CAM_FOV,
+    "near": CAM_NEAR,
+    "far": CAM_FAR,
+    "cx": CAM_CX,
+    "cy": CAM_CY,
+    "fx": CAM_FX,
+    "fy": CAM_FY,
+}
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="configs/server_view_generator.yaml")
+    args = parser.parse_args()
+    vg = ViewGenerator(args.config, CAMERA_PARAMS)
+    vg.run()