add inference

core/dataset.py

@@ -27,7 +27,7 @@ class NBVReconstructionDataset(BaseDataset):
 
         self.pts_num = config["pts_num"]
         self.type = config["type"]
-        self.cache = config["cache"]
+        self.cache = config.get("cache")
         if self.type == namespace.Mode.TEST:
             self.model_dir = config["model_dir"]
             self.filter_degree = config["filter_degree"]
@@ -105,7 +105,10 @@ class NBVReconstructionDataset(BaseDataset):
             nR_to_world_pose = cam_info["cam_to_world_R"] 
             n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), n_to_world_pose)
             nR_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), nR_to_world_pose)
-            cached_data = self.load_from_cache(scene_name, first_frame_idx, frame_idx)
+
+            cached_data = None
+            if self.cache:
+                cached_data = self.load_from_cache(scene_name, first_frame_idx, frame_idx)
             
             if cached_data is None:
                 depth_L, depth_R = DataLoadUtil.load_depth(view_path, cam_info['near_plane'], cam_info['far_plane'], binocular=True)
@@ -116,7 +119,8 @@ class NBVReconstructionDataset(BaseDataset):
                 point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536)
                 overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R)
                 downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(overlap_points, self.pts_num)
-                self.save_to_cache(scene_name, first_frame_idx, frame_idx, downsampled_target_point_cloud)
+                if self.cache:
+                    self.save_to_cache(scene_name, first_frame_idx, frame_idx, downsampled_target_point_cloud)
             else:
                 downsampled_target_point_cloud = cached_data
                 

core/evaluation.py

@@ -1,43 +1,14 @@
 import torch
-import os
-import json
 import numpy as np
-import subprocess
-import tempfile
-from utils.data_load import DataLoadUtil
 from utils.reconstruction import ReconstructionUtil
 from utils.pose import PoseUtil
 from utils.pts import PtsUtil
+from utils.render import RenderUtil
 import PytorchBoot.stereotype as stereotype
 import PytorchBoot.namespace as namespace
 from PytorchBoot.utils.log_util import Log
 
-def render_pts(cam_pose, scene_path,script_path, model_points_normals, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None):
-    nO_to_world_pose = cam_pose.cpu().numpy()  @ nO_to_nL_pose
-    nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose)
-    
-    
-    with tempfile.TemporaryDirectory() as temp_dir:
-        params = {
-            "cam_pose": nO_to_world_pose.tolist(),
-            "scene_path": scene_path
-        }
-        params_data_path = os.path.join(temp_dir, "params.json")
-        with open(params_data_path, 'w') as f:
-            json.dump(params, f)
-        result = subprocess.run([
-            'blender', '-b', '-P', script_path, '--', temp_dir
-        ], capture_output=True, text=True)
-        if result.returncode != 0:
-            print("Blender script failed:")
-            print(result.stderr)
-            return None
-        path = os.path.join(temp_dir, "tmp")
-        
-        point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True)
-        cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
-        sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree)
-        return sampled_point_cloud
+
 
 @stereotype.evaluation_method("pose_diff")
 class PoseDiff:
@@ -110,7 +81,7 @@ class ConverageRateIncrease:
                 filter_degree = filter_degree_list[idx]
                 nO_to_nL_pose = nO_to_nL_pose_batch[idx]
                 try:
-                    new_pts = render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose)
+                    new_pts, _ = RenderUtil.render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose)
                     pred_cr = self.compute_coverage_rate(scanned_view_pts, new_pts, down_sampled_model_pts, threshold=voxel_threshold)
                 except Exception as e:
                     Log.warning(f"Error in scene {scene_path}, {e}")

core/pipeline.py

@@ -5,7 +5,7 @@ import PytorchBoot.stereotype as stereotype
 from PytorchBoot.factory.component_factory import ComponentFactory
 from PytorchBoot.utils import Log
 
-@stereotype.pipeline("nbv_reconstruction_pipeline", comment="should be tested")
+@stereotype.pipeline("nbv_reconstruction_pipeline")
 class NBVReconstructionPipeline(nn.Module):
     def __init__(self, config):
         super(NBVReconstructionPipeline, self).__init__()
@@ -67,14 +67,13 @@ class NBVReconstructionPipeline(nn.Module):
     def get_seq_feat(self, data):
         scanned_pts_batch = data['scanned_pts']
         scanned_n_to_1_pose_9d_batch = data['scanned_n_to_1_pose_9d']
-        best_to_1_pose_9d_batch = data["best_to_1_pose_9d"]
         pts_feat_seq_list = []
         pose_feat_seq_list = []
-        
+        device = next(self.parameters()).device
         for scanned_pts,scanned_n_to_1_pose_9d in zip(scanned_pts_batch,scanned_n_to_1_pose_9d_batch):
- 
-            scanned_pts = scanned_pts.to(best_to_1_pose_9d_batch.device)
-            scanned_n_to_1_pose_9d = scanned_n_to_1_pose_9d.to(best_to_1_pose_9d_batch.device)
+            
+            scanned_pts = scanned_pts.to(device)
+            scanned_n_to_1_pose_9d = scanned_n_to_1_pose_9d.to(device)
             pts_feat_seq_list.append(self.pts_encoder.encode_points(scanned_pts))
             pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_1_pose_9d))
         

core/seq_dataset.py

@@ -51,7 +51,7 @@ class SeqNBVReconstructionDataset(BaseDataset):
                 "first_frame": first_frame,
                 "max_coverage_rate": max_coverage_rate
             })
-        return datalist
+        return datalist[5:]
 
     def __getitem__(self, index):
         data_item_info = self.datalist[index]
@@ -85,6 +85,7 @@ class SeqNBVReconstructionDataset(BaseDataset):
         voxel_threshold = diag*0.02
         first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose)
         scene_path = os.path.join(self.root_dir, scene_name)
+        model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name)
         data_item = {
             "first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32),
             "first_to_first_9d": np.asarray([first_to_first_9d],dtype=np.float32),
@@ -92,10 +93,11 @@ class SeqNBVReconstructionDataset(BaseDataset):
             "max_coverage_rate": max_coverage_rate,
             "voxel_threshold": voxel_threshold,
             "filter_degree": self.filter_degree,
-            "first_frame_to_world": first_frame_to_world,
-            "first_O_to_first_L_pose": first_O_to_first_L_pose,
+            "first_frame_to_world": np.asarray(first_frame_to_world, dtype=np.float32),
+            "O_to_L_pose": first_O_to_first_L_pose,
             "first_frame_coverage": first_frame_coverage,
-            "scene_path": scene_path
+            "scene_path": scene_path,
+            "model_points_normals": model_points_normals,
         }
         return data_item