3D Visualization of FloorPlan

data_manager/data_manager.py

@@ -62,12 +62,13 @@ def load_data(self):
             # ! List of DepthGT maps
             self.list_depth_maps = [os.path.join(self.mp3d_fpe_scene_dir, f'depth/tiff/{f}.tiff') for f in self.list_kf]
 
-            # ! Load GT floor plan & point cloud
-            self.room_corners, self.axis_corners, self.list_kf_per_room = \
-                self.load_fp_gt(os.path.join(self.mp3d_fpe_scene_dir, 'label.json'))
+            if self.cfg['data.use_gt_poses']:
+                # ! Load GT floor plan & point cloud
+                self.room_corners, self.axis_corners, self.list_kf_per_room = \
+                    self.load_fp_gt(os.path.join(self.mp3d_fpe_scene_dir, 'label.json'))
 
-            # NOTE: pcl_gt is (N, 3) and z-axis is the height
-            self.pcl_gt = read_ply(os.path.join(self.mp3d_fpe_scene_dir, 'pcl.ply'))
+                # NOTE: pcl_gt is (N, 3) and z-axis is the height
+                self.pcl_gt = read_ply(os.path.join(self.mp3d_fpe_scene_dir, 'pcl.ply'))
 
             self.cam = Sphere(shape=self.cfg['data.image_resolution'])
         except:
@@ -88,18 +89,19 @@ def load_camera_poses(self):
         self.poses_est = np.stack(
             list(read_trajectory(estimated_poses_file).values()))
 
-        # ! Loading GT camera poses
-        gt_poses_file = os.path.join(
-            self.mp3d_fpe_scene_dir,
-            'frm_ref.txt')
+        if self.cfg['data.use_gt_poses']:
+            # ! Loading GT camera poses
+            gt_poses_file = os.path.join(
+                self.mp3d_fpe_scene_dir,
+                'frm_ref.txt')
 
-        assert os.path.isfile(
-            gt_poses_file
-        ), f'Cam pose file {gt_poses_file} does not exist'
+            assert os.path.isfile(
+                gt_poses_file
+            ), f'Cam pose file {gt_poses_file} does not exist'
 
-        idx = np.array(self.list_kf)-1
-        self.poses_gt = np.stack(
-            list(read_trajectory(gt_poses_file).values()))[idx, :, :]
+            idx = np.array(self.list_kf)-1
+            self.poses_gt = np.stack(
+                list(read_trajectory(gt_poses_file).values()))[idx, :, :]
 
     def load_fp_gt(self, fn):
         """
@@ -149,7 +151,7 @@ def get_list_ly(self, cam_ref=CAM_REF.CC):
         for idx_kf in tqdm(self.list_kf, desc="Loading Layouts..."):
             idx = self.list_kf.index(idx_kf)
 
-            if self.cfg['data.use_gt_poses']:
+            if not self.cfg['data.use_gt_poses']:
                 pose = CamPose(self, pose=self.poses_est[idx])
             else:
                 pose = CamPose(self, pose=self.poses_gt[idx])
@@ -185,7 +187,8 @@ def get_list_ly(self, cam_ref=CAM_REF.CC):
             _, s, _ = np.linalg.svd(cov/(pcl.size - 1))
 
             ly = Layout(self)
-            ly.pose_gt = CamPose(self, pose=self.poses_gt[idx])
+            if self.cfg['data.use_gt_poses']:
+                ly.pose_gt = CamPose(self, pose=self.poses_gt[idx])
             ly.bearings = bearings
             ly.boundary = pcl
             ly.pose = pose

direct_floor_plan_estimation/scale_recover/scale_recover.py

@@ -144,11 +144,12 @@ def estimate_vo_and_gt_scale(self):
         """
 
         self.estimate_vo_scale()
-        self.gt_scale = self.gt_scale_recover.estimate(self)
+        if self.dt.cfg['data.use_gt_poses']:
+            self.gt_scale = self.gt_scale_recover.estimate(self)
 
-        # * We are assuming that by estimating vo and gt scale is because we want to
-        # * apply GT scale
-        [ly.apply_gt_scale(self.gt_scale) for ly in self.dt.list_ly]
+            # * We are assuming that by estimating vo and gt scale is because we want to
+            # * apply GT scale
+            [ly.apply_gt_scale(self.gt_scale) for ly in self.dt.list_ly]
 
         return self.gt_scale, self.vo_scale
 

main_eval_scene.py

@@ -12,7 +12,7 @@
 from utils.visualization.room_utils import plot_floor_plan, plot_all_planes, plot_planes_rooms_patches
 from utils.eval_utils import evaluate_corners_pr, evaluate_rooms_pr
 from utils.io import read_scene_list
-from utils.eval_utils import evaluate_scene, dump_images, dump_result
+from utils.eval_utils import evaluate_scene, dump_images, dump_result, dump_fpe_geometry
 
 
 def main(opt):
@@ -22,6 +22,7 @@ def main(opt):
 
     cfg["data.scene"] = opt.scene_name.split("_")[0]
     cfg["data.scene_version"] = opt.scene_name.split("_")[1]
+    cfg['data.use_gt_poses'] = True
 
     dt = DataManager(cfg)
     fpe = DirectFloorPlanEstimation(dt)
@@ -33,39 +34,45 @@ def main(opt):
     fpe.eval_room_overlapping()
     fpe.masking_ocg_map()
 
-    points_gt = fpe.dt.pcl_gt      # (3, N)
+    if cfg['data.use_gt_poses']:
+        points_gt = fpe.dt.pcl_gt      # (3, N)
 
     room_corner_list = fpe.compute_room_shape_all()
     image_room_id = plot_all_rooms_by_patches(fpe)
     image_final_fp = plot_floor_plan(room_corner_list, fpe.global_ocg_patch)
     room_corner_list = [x.T for x in room_corner_list]  # Make it (N, 2)
-    result_dict, images_dict = evaluate_scene(
-        room_corner_list,
-        fpe.dt.room_corners,
-        points_gt,
-        axis_corners=fpe.dt.axis_corners
-    )
-    result_dict['scene'] = cfg['data.scene']
-    images_dict['scene'] = cfg['data.scene']
-    images_dict['room_id'] = image_room_id
-    images_dict['final_fp'] = image_final_fp
+    if cfg['data.use_gt_poses']:
+        result_dict, images_dict = evaluate_scene(
+            room_corner_list,
+            fpe.dt.room_corners,
+            points_gt,
+            axis_corners=fpe.dt.axis_corners
+        )
+        result_dict['scene'] = cfg['data.scene']
+        images_dict['scene'] = cfg['data.scene']
+        images_dict['room_id'] = image_room_id
+        images_dict['final_fp'] = image_final_fp
 
     # Saving the results
     results_dir = os.path.join(output_dir, f"{cfg['data.scene']}_{cfg['data.scene_version']}")
     os.makedirs(results_dir, exist_ok=True)
 
-    # GT data for references
-    dt.save_gt_rooms(results_dir)
+    if cfg['data.use_gt_poses']:
 
-    # Estimated VO-SCALE and density 2d function
-    fpe.scale_recover.save_estimation(results_dir)
+        # GT data for references
+        dt.save_gt_rooms(results_dir)
 
-    #  Estimated results
-    dump_images(images_dict, results_dir)
+        # Estimated VO-SCALE and density 2d function
+        fpe.scale_recover.save_estimation(results_dir)
 
-    # writing results
-    dump_result([result_dict], output_dir)
+        #  Estimated results
+        dump_images(images_dict, results_dir)
 
+        # writing results
+        dump_result([result_dict], output_dir)
+
+    #writing fpe geomtery
+    dump_fpe_geometry(room_corner_list, results_dir)
 
 def get_passed_args():
     parser = argparse.ArgumentParser()

main_geometry_visualization.py

@@ -0,0 +1,15 @@
+import os
+import argparse
+import open3d as o3d
+from utils.visualization.open3d_utils import visualize_geometry
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--pkl', required=True, help='path to load pickle file.')
+    parser.add_argument('--texture_dir', required=False, default=None, help='path to load config file.')
+
+    args = parser.parse_args()
+
+    dfpe_layout = visualize_geometry(pkl_geometry_path=args.pkl)
+    o3d.io.write_triangle_mesh(os.path.join(os.path.dirname(args.pkl), "3D_layout.obj"), dfpe_layout)

requirements.txt

@@ -20,3 +20,5 @@ googledrivedownloader
 scikit-image
 hydra-core
 wandb==0.15.3
+opencv-python
+open3d

utils/eval_utils.py

@@ -92,6 +92,11 @@ def dump_images(images_dict, save_dir):
         plt.imsave(os.path.join(save_dir, f'{scene}_{key}_gt.png'), images_dict[key]['image_gt'])
         plt.imsave(os.path.join(save_dir, f'{scene}_{key}_pred.png'), images_dict[key]['image_pred'])
 
+def dump_fpe_geometry(fpe_geometry_list, save_dir):
+    import pickle
+    fpe_geometry = {'corners_list': fpe_geometry_list}
+    with open(os.path.join(save_dir, "fpe_corners.pkl"), 'wb') as f:
+        pickle.dump(fpe_geometry, f)
 
 def dump_result(result_list, save_dir):
     # Dump result to a csv file

utils/visualization/open3d_utils.py

@@ -0,0 +1,54 @@
+import pickle
+import numpy as np
+import open3d as o3d
+
+def visualize_geometry(pkl_geometry_path):
+    try:
+        fpe = pickle.load(open(pkl_geometry_path, 'rb')) 
+    except:
+        print("No pickle file found")
+
+    all_rooms_corners = fpe["corners_list"]
+
+    R = 0.8
+    cam_height = 1.5
+    rooms_mesh = []
+    layout_mesh = o3d.geometry.TriangleMesh()
+    room_mesh_subdivided = o3d.geometry.TriangleMesh()
+
+    pts = []
+    li = []
+    triangles = []
+    for room_corners in all_rooms_corners:
+        x_list = room_corners[:, 0].tolist()
+        z_list = room_corners[:, 1].tolist()
+        k = len(pts)
+        n = len(x_list)
+        centroid_up = [(np.sum(x_list)/len(x_list), cam_height, np.sum(z_list)/len(z_list))]
+        centroid_down = [(np.sum(x_list)/len(x_list), -R * cam_height, np.sum(z_list)/len(z_list))]
+        for x,z in zip(x_list, z_list):
+            # Despite of what's on the literature, cam_height * R gives height that looks better when textured
+            pts += [(x, cam_height, z), (x, - cam_height * R, z)]
+            for i in range(0, 2*len(x_list), 2):
+                triangles.append([k+(i+0)%(2*n), k+(i+1)%(2*n), k+(i+2)%(2*n)])
+                triangles.append([k+(i+1)%(2*n), k+(i+3)%(2*n), k+(i+2)%(2*n)])
+                triangles.append([k+(i+0)%(2*n), k+(i+2)%(2*n), 2*len(x_list)])
+
+                li.append([k+i, k+(i+2)%(2*n)])
+                li.append([k+i, k+(i+1)])
+                li.append([k+i+1, k+(i+3)%(2*n)])
+                li.append([k+i+1, 2*len(x_list)+1])
+                li.append([k+i, 2*len(x_list)])
+
+        pts = pts+centroid_up+centroid_down
+
+        room_mesh = o3d.geometry.TriangleMesh()
+        room_mesh.vertices = o3d.utility.Vector3dVector(pts)
+        room_mesh.triangles = o3d.utility.Vector3iVector(triangles)
+
+        rooms_mesh.append(room_mesh)
+        room_mesh_subdivided = room_mesh.subdivide_midpoint(3)
+        layout_mesh += room_mesh_subdivided
+
+    o3d.visualization.draw(room_mesh_subdivided)
+    return room_mesh_subdivided