添加状态分类和液面分割

angle_base_seg/angle_main.py

@@ -4,7 +4,7 @@ import numpy as np
 import os
 
 # ------------------ 模型与路径配置 ------------------
-MODEL_PATH = '../ultralytics_yolo11-main/runs/train/exp4/weights/best.pt'
+MODEL_PATH = '../ultralytics_yolo11-main/runs/train/seg_j/exp/weights/best.pt'
 OUTPUT_DIR = '../test_image'
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
@@ -34,7 +34,7 @@ def detect_jaw_angle(image_path, mode='show'):
 
     # 创建掩码并检测
     composite_mask = np.zeros((h, w), dtype=np.uint8)
-    results = model(image_path, imgsz=1280, conf=0.5)
+    results = model(image_path, imgsz=640, conf=0.5)
 
     jaws = []
     for r in results:
@@ -143,13 +143,13 @@ def detect_jaw_angle(image_path, mode='show'):
 
 # ------------------ 主函数 ------------------
 if __name__ == '__main__':
-    # ✅ 设置输入图像路径
-    image_path = '/test_image/1.png'  # ← 修改为你自己的图片路径
+    # ✅ 设置输入图像路
+    image_path = r"/home/hx/yolo/output_masks/2.jpg"  # ← 修改为你自己的图片路径
 
     # ✅ 模式选择：
     #   mode='show': 保存可视化图像
     #   mode='silent': 只返回角度
-    mode = 'silent'  # 或 'silent'
+    mode = 'show'  # 或 'silent'
 
     print(f"🔍 正在处理图像: {image_path}")
     angle = detect_jaw_angle(image_path, mode=mode)