更新 'GFS-Client/save_result.py'

GFS-Client/save_result.py

@@ -121,19 +121,19 @@ def post_process_in_2d(instances_with_pixel_boxes, x_m_per_px, y_m_per_px, iou_t
         metric_height = px_height * y_m_per_px
         extent_2d = [metric_width, metric_height]
         
-        if is_box_dimension_plausible_2d(extent_2d, inst['label']):
+        if is_box_dimension_plausible_2d(extent_2d, inst['category']):
             plausible_instances.append(inst)
         else:
-             print(f"  - 过滤掉一个2D尺寸异常的 '{inst['label']}' 实例，尺寸: {[f'{x:.2f}' for x in extent_2d]}")
+             print(f"  - 过滤掉一个2D尺寸异常的 '{inst['category']}' 实例，尺寸: {[f'{x:.2f}' for x in extent_2d]}")
 
     if not plausible_instances:
         return []
 
     # 2. 按类别分组进行后处理
     final_instances = []
-    plausible_instances.sort(key=lambda x: x['label'])
+    plausible_instances.sort(key=lambda x: x['category'])
     
-    for class_name, group in groupby(plausible_instances, key=lambda x: x['label']):
+    for class_name, group in groupby(plausible_instances, key=lambda x: x['category']):
         class_instances = list(group)
         
         # --- SPECIAL MERGING LOGIC FOR BEDS ---
@@ -173,18 +173,34 @@ def post_process_in_2d(instances_with_pixel_boxes, x_m_per_px, y_m_per_px, iou_t
             for group_indices in groups:
                 instances_in_group = [class_instances[i] for i in group_indices]
                 
-                # Create the merged bounding box
-                min_x = min(inst['bbox_2d_pixels'][0] for inst in instances_in_group)
-                min_y = min(inst['bbox_2d_pixels'][1] for inst in instances_in_group)
-                max_x = max(inst['bbox_2d_pixels'][2] for inst in instances_in_group)
-                max_y = max(inst['bbox_2d_pixels'][3] for inst in instances_in_group)
+                # Create the merged 2D bounding box
+                min_x_2d = min(inst['bbox_2d_pixels'][0] for inst in instances_in_group)
+                min_y_2d = min(inst['bbox_2d_pixels'][1] for inst in instances_in_group)
+                max_x_2d = max(inst['bbox_2d_pixels'][2] for inst in instances_in_group)
+                max_y_2d = max(inst['bbox_2d_pixels'][3] for inst in instances_in_group)
                 
+                # --- Create the merged 3D bounding box ---
+                all_3d_corners = np.vstack([inst['bbox'] for inst in instances_in_group])
+                min_3d = np.min(all_3d_corners, axis=0)
+                max_3d = np.max(all_3d_corners, axis=0)
+                merged_3d_bbox = [
+                    [min_3d[0], min_3d[1], min_3d[2]],
+                    [max_3d[0], min_3d[1], min_3d[2]],
+                    [min_3d[0], max_3d[1], min_3d[2]],
+                    [max_3d[0], max_3d[1], min_3d[2]],
+                    [min_3d[0], min_3d[1], max_3d[2]],
+                    [max_3d[0], min_3d[1], max_3d[2]],
+                    [min_3d[0], max_3d[1], max_3d[2]],
+                    [max_3d[0], max_3d[1], max_3d[2]],
+                ]
+
                 # Aggregate score and find a representative instance for metadata
                 total_score = sum(inst['score'] for inst in instances_in_group)
                 representative_instance = max(instances_in_group, key=lambda x: x['score'])
                 
                 new_instance = representative_instance.copy()
-                new_instance['bbox_2d_pixels'] = [min_x, min_y, max_x, max_y]
+                new_instance['bbox_2d_pixels'] = [min_x_2d, min_y_2d, max_x_2d, max_y_2d]
+                new_instance['bbox'] = merged_3d_bbox # Assign the new merged 3D bbox
                 new_instance['score'] = total_score
                 merged_instances.append(new_instance)
             
@@ -242,13 +258,16 @@ def build_floor_transform_matrix(j_info: dict, floor_id: int):
     return np.linalg.inv(tab_array).tolist(), res_width, res_height
 
 
-def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id, output_image_path, output_json_path):
+def process_and_draw_bboxes(picture_name, floor_path, raw_bbox_data, floor_id, output_image_path, output_json_path, output_3d_json_path):
     try:
         img = cv2.imread(picture_name)
         if img is None: raise FileNotFoundError(f"无法加载背景图片: {picture_name}")
         
         with open(floor_path, 'r', encoding='utf-8') as f: j_info = json.load(f)
-        with open(instances_path, 'r', encoding='utf-8') as f: raw_bbox_data = json.load(f)
+
+        if not raw_bbox_data:
+            print("警告: 未提供任何原始3D包围盒数据。")
+            return None, None
 
         matrix, res_w, res_h = build_floor_transform_matrix(j_info, floor_id)
         if res_w is None: raise ValueError(f"未在 {floor_path} 中找到 ID 为 {floor_id} 的楼层信息。")
@@ -261,11 +280,11 @@ def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id,
 
         instances_with_pixel_boxes = []
         for item in raw_bbox_data:
-            corners = item.get("corners", [])
-            if len(corners) < 4: continue
+            corners = item.get("bbox", [])
+            if len(corners) < 8: continue
             
             points_2d = []
-            for i in range(4):
+            for i in range(8):
                 norm_pt = M @ np.array([corners[i][0], corners[i][1], 1.0])
                 points_2d.append([int(norm_pt[0] * res_w), int(norm_pt[1] * res_h)])
             
@@ -279,10 +298,10 @@ def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id,
         print("2D后处理完成。")
         
         img_height, img_width, _ = img.shape
-        shapes = []
+        shapes_2d = []
         for item in filtered_bbox_data:
             min_x, min_y, max_x, max_y = item['bbox_2d_pixels']
-            category = item["label"]
+            category = item["category"]
             color_rgb = item["color"]
             color_bgr = (color_rgb[2], color_rgb[1], color_rgb[0])
 
@@ -296,7 +315,7 @@ def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id,
             bbox_poly = [min_x, min_y, max_x, min_y, max_x, max_y, min_x, max_y]
             class_info = CLASS_MAPPING.get(category, {'id': '-1', 'name': '未知'})
 
-            shapes.append({
+            shapes_2d.append({
                 "bbox": bbox_poly,
                 "category": category,
                 "color": color_rgb,
@@ -304,20 +323,41 @@ def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id,
                 "name": class_info['name']
             })
 
-        output_json_data = {
-            "shapes": shapes,
+        # --- Save 2D Results ---
+        output_2d_json_data = {
+            "shapes": shapes_2d,
             "imageHeight": img_height,
             "imagePath": os.path.basename(picture_name),
             "imageWidth": img_width,
             "version": "4Dage_Furniture_Detection_0.0.1"
         }
-
         os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
         os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
         cv2.imwrite(output_image_path, img)
         with open(output_json_path, 'w', encoding='utf-8') as f:
-            json.dump(output_json_data, f, ensure_ascii=False, indent=4)
+            json.dump(output_2d_json_data, f, ensure_ascii=False, indent=4)
         print(f"\n处理完成！2D结果已保存到: {output_image_path} 和 {output_json_path}")
+        
+        # --- Save Final 3D Results ---
+        shapes_3d = []
+        for item in filtered_bbox_data:
+            shapes_3d.append({
+                "bbox": item['bbox'],
+                "category": item['category'],
+                "color": item['color'],
+                "label": item['label'],
+                "name": item['name'],
+            })
+        
+        output_3d_json_data = {
+            "shapes": shapes_3d,
+            "version": "4Dage_Furniture_Detection_0.0.1_3D_final"
+        }
+        os.makedirs(os.path.dirname(output_3d_json_path), exist_ok=True)
+        with open(output_3d_json_path, 'w', encoding='utf-8') as f:
+            json.dump(output_3d_json_data, f, ensure_ascii=False, indent=4)
+        print(f"对应的3D结果已保存到: {output_3d_json_path}")
+
         return output_json_path, output_image_path
 
     except Exception as e:
@@ -331,7 +371,6 @@ def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id,
 def visualize_point_cloud_segmentation(coords_file, preds_file, classes_to_show='all',
                                        classes_to_ignore=None,
                                        save_pcd_path=None,
-                                       save_3d_json_path=None,
                                        if_save_ply=False,
                                        if_save_vision=False):
     CLASS_NAMES = [
@@ -387,20 +426,20 @@ def visualize_point_cloud_segmentation(coords_file, preds_file, classes_to_show=
             try:
                 aabb = instance_pcd.get_axis_aligned_bounding_box()
                 points_np = np.asarray(instance_pcd.points)
+                class_info = CLASS_MAPPING.get(class_name, {'id': '-1', 'name': '未知'})
                 final_instances_data.append({
-                    "label": class_name, "color": COLOR_MAP[pred_idx].tolist(),
-                    "corners": np.asarray(aabb.get_box_points()).tolist(), "score": len(points_np)
+                    "category": class_name,
+                    "label": class_info['id'],
+                    "name": class_info['name'],
+                    "color": COLOR_MAP[pred_idx].tolist(),
+                    "bbox": np.asarray(aabb.get_box_points()).tolist(),
+                    "score": len(points_np)
                 })
                 all_instance_points.append(points_np)
                 all_instance_colors.append(np.tile(COLOR_MAP[pred_idx] / 255.0, (len(points_np), 1)))
             except RuntimeError: continue
 
     print("\n所有原始实例处理完毕。")
-    if save_3d_json_path:
-        os.makedirs(os.path.dirname(save_3d_json_path), exist_ok=True)
-        with open(save_3d_json_path, 'w', encoding='utf-8') as f:
-            json.dump(final_instances_data, f, ensure_ascii=False, indent=4)
-        print(f"原始3D实例JSON信息已保存至: {save_3d_json_path}")
 
     if if_save_ply and save_pcd_path and all_instance_points:
         instance_pcd = o3d.geometry.PointCloud()
@@ -414,7 +453,7 @@ def visualize_point_cloud_segmentation(coords_file, preds_file, classes_to_show=
                                       colors=o3d.utility.Vector3dVector(COLOR_MAP[predictions] / 255.0))
         o3d.visualization.draw_geometries([pcd], window_name="原始点云", width=1280, height=720)
 
-    return save_3d_json_path
+    return final_instances_data
 
 
 if __name__ == "__main__":
@@ -440,23 +479,28 @@ if __name__ == "__main__":
 
     output_dir = os.path.join(scene_folder, 'result_2d_filtered')
     os.makedirs(output_dir, exist_ok=True)
-    raw_instances3d_json_path = os.path.join(output_dir, 'instances3d_raw.json')
+    
     final_instances2d_json_path = os.path.join(output_dir, 'instances2d_final.json')
+    final_instances3d_json_path = os.path.join(output_dir, 'instances3d_final.json')
     instances_ply_path = os.path.join(output_dir, 'instances_raw.ply')
     segment_onfloor_png_path = os.path.join(output_dir, 'segment_onfloor_final.png')
 
-    saved_3d_json = visualize_point_cloud_segmentation(
+    raw_3d_instances = visualize_point_cloud_segmentation(
         coords_file=coords_file, preds_file=preds_file,
         classes_to_ignore=['curtain', 'bookshelf', 'floor', 'wall', 'sink', 'toilet', 'bathtub', 'shower curtain', 'picture'],
-        save_3d_json_path=raw_instances3d_json_path, save_pcd_path=instances_ply_path, if_save_ply=False
+        save_pcd_path=instances_ply_path, if_save_ply=False
     )
 
-    if saved_3d_json and all(os.path.exists(f) for f in [floor_plan_image, scene_info_json]):
+    if raw_3d_instances and all(os.path.exists(f) for f in [floor_plan_image, scene_info_json]):
         print("\n--- 开始进行2D投影和后处理 ---")
         process_and_draw_bboxes(
-            picture_name=floor_plan_image, floor_path=scene_info_json,
-            instances_path=saved_3d_json, floor_id=0,
-            output_image_path=segment_onfloor_png_path, output_json_path=final_instances2d_json_path
+            picture_name=floor_plan_image, 
+            floor_path=scene_info_json,
+            raw_bbox_data=raw_3d_instances, 
+            floor_id=0,
+            output_image_path=segment_onfloor_png_path, 
+            output_json_path=final_instances2d_json_path,
+            output_3d_json_path=final_instances3d_json_path
         )
     else:
-        print("\nSkipping 2D projection due to missing files.")
+        print("\nSkipping 2D projection due to missing files or no raw instances detected.")