success

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')