|
@@ -8,6 +8,29 @@ from itertools import groupby
|
|
|
import argparse
|
|
import argparse
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+CLASS_MAPPING = {
|
|
|
|
|
+ 'refrigerator': {'id': '0', 'name': '冰箱'},
|
|
|
|
|
+ 'desk': {'id': '1', 'name': '书桌'},
|
|
|
|
|
+ 'curtain': {'id': '2', 'name': '窗帘'},
|
|
|
|
|
+ 'sofa': {'id': '3', 'name': '沙发'},
|
|
|
|
|
+ 'bookshelf': {'id': '4', 'name': '书架'},
|
|
|
|
|
+ 'bed': {'id': '5', 'name': '床'},
|
|
|
|
|
+ 'table': {'id': '6', 'name': '桌子'},
|
|
|
|
|
+ 'window': {'id': '7', 'name': '窗户'},
|
|
|
|
|
+ 'cabinet': {'id': '8', 'name': '橱柜'},
|
|
|
|
|
+ 'door': {'id': '9', 'name': '门'},
|
|
|
|
|
+ 'chair': {'id': '10', 'name': '椅子'},
|
|
|
|
|
+ 'floor': {'id': '11', 'name': '地板'},
|
|
|
|
|
+ 'wall': {'id': '12', 'name': '墙'},
|
|
|
|
|
+ 'sink': {'id': '13', 'name': '水槽'},
|
|
|
|
|
+ 'toilet': {'id': '14', 'name': '马桶'},
|
|
|
|
|
+ 'bathtub': {'id': '15', 'name': '浴缸'},
|
|
|
|
|
+ 'shower curtain': {'id': '16', 'name': '浴帘'},
|
|
|
|
|
+ 'picture': {'id': '17', 'name': '画'},
|
|
|
|
|
+ 'counter': {'id': '18', 'name': '柜台'},
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
# ==============================================================================
|
|
# ==============================================================================
|
|
|
# 优化算法辅助函数
|
|
# 优化算法辅助函数
|
|
|
# ==============================================================================
|
|
# ==============================================================================
|
|
@@ -255,26 +278,45 @@ def process_and_draw_bboxes(picture_name, floor_path, instances_path, floor_id,
|
|
|
filtered_bbox_data = post_process_in_2d(instances_with_pixel_boxes, x_m_per_px, y_m_per_px)
|
|
filtered_bbox_data = post_process_in_2d(instances_with_pixel_boxes, x_m_per_px, y_m_per_px)
|
|
|
print("2D后处理完成。")
|
|
print("2D后处理完成。")
|
|
|
|
|
|
|
|
- instances_2d_data = []
|
|
|
|
|
|
|
+ img_height, img_width, _ = img.shape
|
|
|
|
|
+ shapes = []
|
|
|
for item in filtered_bbox_data:
|
|
for item in filtered_bbox_data:
|
|
|
min_x, min_y, max_x, max_y = item['bbox_2d_pixels']
|
|
min_x, min_y, max_x, max_y = item['bbox_2d_pixels']
|
|
|
- label = item["label"]
|
|
|
|
|
- color_bgr = (item["color"][2], item["color"][1], item["color"][0])
|
|
|
|
|
-
|
|
|
|
|
- instances_2d_data.append({"label": label, "color": item["color"], "bbox_2d": item['bbox_2d_pixels']})
|
|
|
|
|
|
|
+ category = item["label"]
|
|
|
|
|
+ color_rgb = item["color"]
|
|
|
|
|
+ color_bgr = (color_rgb[2], color_rgb[1], color_rgb[0])
|
|
|
|
|
|
|
|
cv2.rectangle(img, (min_x, min_y), (max_x, max_y), color_bgr, 2)
|
|
cv2.rectangle(img, (min_x, min_y), (max_x, max_y), color_bgr, 2)
|
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
|
- (text_w, text_h), _ = cv2.getTextSize(label, font, 0.5, 1)
|
|
|
|
|
|
|
+ (text_w, text_h), _ = cv2.getTextSize(category, font, 0.5, 1)
|
|
|
label_y = min_y - 10 if min_y - 10 > text_h else min_y + text_h + 10
|
|
label_y = min_y - 10 if min_y - 10 > text_h else min_y + text_h + 10
|
|
|
cv2.rectangle(img, (min_x, label_y - text_h - 5), (min_x + text_w, label_y + 5), color_bgr, -1)
|
|
cv2.rectangle(img, (min_x, label_y - text_h - 5), (min_x + text_w, label_y + 5), color_bgr, -1)
|
|
|
- cv2.putText(img, label, (min_x, label_y), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
|
|
|
|
|
|
|
+ cv2.putText(img, category, (min_x, label_y), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
|
|
|
|
|
+
|
|
|
|
|
+ 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({
|
|
|
|
|
+ "bbox": bbox_poly,
|
|
|
|
|
+ "category": category,
|
|
|
|
|
+ "color": color_rgb,
|
|
|
|
|
+ "label": class_info['id'],
|
|
|
|
|
+ "name": class_info['name']
|
|
|
|
|
+ })
|
|
|
|
|
+
|
|
|
|
|
+ output_json_data = {
|
|
|
|
|
+ "shapes": shapes,
|
|
|
|
|
+ "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_image_path), exist_ok=True)
|
|
|
os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
|
|
os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
|
|
|
cv2.imwrite(output_image_path, img)
|
|
cv2.imwrite(output_image_path, img)
|
|
|
with open(output_json_path, 'w', encoding='utf-8') as f:
|
|
with open(output_json_path, 'w', encoding='utf-8') as f:
|
|
|
- json.dump(instances_2d_data, f, indent=4)
|
|
|
|
|
|
|
+ json.dump(output_json_data, f, ensure_ascii=False, indent=4)
|
|
|
print(f"\n处理完成!2D结果已保存到: {output_image_path} 和 {output_json_path}")
|
|
print(f"\n处理完成!2D结果已保存到: {output_image_path} 和 {output_json_path}")
|
|
|
return output_json_path, output_image_path
|
|
return output_json_path, output_image_path
|
|
|
|
|
|