修改了识别逻辑

predict/predict_yolo11seg.py

@@ -2,7 +2,7 @@ import logging
 import shutil
 import zipfile
 from os.path import exists
-
+from pathlib import Path
 import torch
 import gc
 import os
@@ -385,7 +385,7 @@ class YOLOSegmentationInference:
         return output_image
 
     def save_results(self, result: InferenceResult, output_dir: str, save_mask: bool = False,
-                     save_label: bool = False, result_save: [] = None) -> None:
+                     save_label: bool = False, result_save: List[Dict] = None) -> None:
         """
         保存推理结果
 
@@ -510,43 +510,125 @@ class YOLOSegmentationInference:
     def process_single_image(self, image_path: str, output_dir: Optional[str] = None,
                             conf_threshold: float = 0.25, iou_threshold: float = 0.5,
                             save_mask: bool = False, save_label: bool = False, show: bool = True,
-                             result_save: [] = None) -> InferenceResult:
+                            result_save: Optional[List] = None):
         """
-        处理单张图片
+        处理单张图片（使用YOLO分割模型）
         
         Args:
-            image_path: 图片路径
+            image_path: 输入图片路径
-            output_dir: 输出目录，如果为None则不保存
+            output_dir: 输出目录（如果为None则不保存）
             conf_threshold: 置信度阈值
-            iou_threshold: IOU阈值
+            iou_threshold: NMS IoU阈值
-            save_mask: 是否保存单独的掩码文件
+            save_mask: 是否保存单独的掩码图像（PNG格式）
-            save_label: 是否保存YOLO格式的标签文件
+            save_label: 是否保存YOLO格式标签文件（txt，含多边形或边界框）
-            show: 是否显示结果
+            show: 是否显示结果图像（按任意键关闭）
+            result_save: 可选的外部列表，用于收集本次处理的信息（如图像路径、保存的文件等）
 
         Returns:
-            推理结果
+            Ultralytics 的 Results 对象（包含原始图像、边界框、掩码等信息）
         """
-        # 执行推理
+        # 1. 确保模型已加载
-        result = self.perform_inference(image_path, conf_threshold, iou_threshold)
+        if not hasattr(self, 'model') or self.model is None:
+            raise ValueError("模型未加载，请先初始化模型（self.model = YOLO('...')）")
 
-        # 绘制结果
+        # 2. 执行推理
-        if result.masks is not None and len(result.masks) > 0:
+        results = self.model.predict(
-            self.draw_results(result, conf_threshold)
+            source=image_path,
+            conf=conf_threshold,
+            iou=iou_threshold,
+            imgsz=640,                # 可根据需要调整，或作为参数传入
+            save=False,                # 手动保存，不依赖自动保存
+            save_txt=False,            # 手动处理标签
+            retina_masks=True,         # 获取与原图同分辨率的高精度掩码
+        )
+        # 单张图片，取第一个结果
+        result = results[0]
 
-        # 保存结果
+        # 3. 准备输出目录
-        if output_dir is not None:
+        if output_dir:
-            self.save_results(result, output_dir, save_mask, save_label, result_save)
+            output_path = Path(output_dir)
+            output_path.mkdir(parents=True, exist_ok=True)
+            base_name = Path(image_path).stem
+        else:
+            output_path = None
+            base_name = None
 
-        # # 显示结果
+        # 4. 绘制标注图像（使用 result.plot()）
+        annotated_img = result.plot()  # 返回BGR numpy数组
+
+        # 5. 保存标注图像到输出目录（如果指定）
+        if output_path:
+            img_save_path = output_path / f"{base_name}_annotated.jpg"
+            cv2.imwrite(str(img_save_path), annotated_img)
+        else:
+            img_save_path = None
+
+        # 6. 保存单独的掩码文件（如果 save_mask 为 True 且有掩码）
+        mask_save_paths = []
+        if save_mask and result.masks is not None and output_path:
+            masks = result.masks.data.cpu().numpy()  # (N, H, W) 二值掩码
+            for i, mask in enumerate(masks):
+                mask_img = (mask * 255).astype(np.uint8)   # 转换为0-255图像
+                mask_file = output_path / f"{base_name}_mask_{i}.png"
+                cv2.imwrite(str(mask_file), mask_img)
+                mask_save_paths.append(str(mask_file))
+
+        # 7. 保存YOLO格式标签文件（如果 save_label 为 True）
+        label_save_path = None
+        if save_label and output_path:
+            # 优先保存分割多边形标签（masks.xy）
+            if result.masks is not None and hasattr(result.masks, 'xy') and result.masks.xy:
+                label_file = output_path / f"{base_name}.txt"
+                with open(label_file, 'w') as f:
+                    for i, poly in enumerate(result.masks.xy):
+                        cls_id = int(result.boxes.cls[i]) if result.boxes is not None else 0
+                        # 归一化多边形点坐标（除以图像宽高）
+                        h, w = result.orig_shape
+                        normalized = poly / [w, h]
+                        flat_coords = normalized.flatten().tolist()
+                        line = f"{cls_id} " + " ".join([f"{x:.6f}" for x in flat_coords])
+                        f.write(line + "\n")
+                label_save_path = str(label_file)
+            # 如果没有掩码但有边界框，则保存检测标签（YOLO检测格式）
+            elif result.boxes is not None:
+                label_file = output_path / f"{base_name}.txt"
+                with open(label_file, 'w') as f:
+                    boxes = result.boxes.xyxyn.cpu().numpy()   # 归一化边界框 (xyxy)
+                    cls_ids = result.boxes.cls.cpu().numpy().astype(int)
+                    for box, cls_id in zip(boxes, cls_ids):
+                        # 转换为 YOLO 格式 (x_center, y_center, width, height)
+                        x1, y1, x2, y2 = box
+                        x_center = (x1 + x2) / 2
+                        y_center = (y1 + y2) / 2
+                        width = x2 - x1
+                        height = y2 - y1
+                        line = f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
+                        f.write(line + "\n")
+                label_save_path = str(label_file)
+
+        # # 8. 显示结果（如果 show 为 True）
         # if show:
-        #     self.show_results(result)
+        #     cv2.imshow("Segmentation Result", annotated_img)
+        #     cv2.waitKey(0)
+        #     cv2.destroyAllWindows()
 
+        # 9. 如果提供了 result_save 列表，将本次结果信息追加进去
+        if result_save is not None:
+            result_save.append({
+                'image_path': image_path,
+                'annotated_image_path': str(img_save_path) if img_save_path else None,
+                'mask_paths': mask_save_paths,
+                'label_path': label_save_path,
+                'num_objects': len(result.boxes) if result.boxes else 0,
+            })
+
+        # 10. 返回结果（Ultralytics 的 Results 对象，包含所有推理信息）
         return result
     
     def process_single_image_share_dir(self, image_path, user_name, pwd, output_dir: Optional[str] = None,
                              conf_threshold: float = 0.25, iou_threshold: float = 0.5,
                              save_mask: bool = False, save_label: bool = False, show: bool = True,
-                             result_save: [] = None) -> None:
+                             result_save: List[Dict] = None) -> None:
         """
         处理单张图片
 
@@ -585,7 +667,7 @@ class YOLOSegmentationInference:
     def process_image_directory(self, input_dir: str, output_dir: Optional[str] = None,
                                 conf_threshold: float = 0.25, iou_threshold: float = 0.5,
                                 save_mask: bool = False, save_label: bool = False, show: bool = False,
-                                result_save: [] = None) -> List[
+                                result_save: List[Dict] = None) -> List[
         InferenceResult]:
         """
         处理目录中的所有图片
@@ -681,7 +763,7 @@ class YOLOSegmentationInference:
     def process_image_directory_share_dir_circle(self, task_id, current_time, input_dir_list, user_name, pwd, output_dir: Optional[str] = None,
                                 conf_threshold: float = 0.25, iou_threshold: float = 0.5,
                                 save_mask: bool = False, save_label: bool = False, show: bool = False,
-                                result_save: [] = None) -> None:
+                                result_save: List[Dict] = None) -> None:
         for input_dir in input_dir_list :
             self.process_image_directory_share_dir(task_id,current_time,input_dir,user_name,pwd,output_dir,conf_threshold,iou_threshold,save_mask,save_label,show,result_save)
         del_file_shutil(output_dir)
@@ -689,7 +771,7 @@ class YOLOSegmentationInference:
     def process_image_directory_share_dir(self, task_id, current_time, input_dir, user_name, pwd, output_dir: Optional[str] = None,
                                 conf_threshold: float = 0.25, iou_threshold: float = 0.5,
                                 save_mask: bool = False, save_label: bool = False, show: bool = False,
-                                result_save: [] = None) -> None:
+                                result_save: List[Dict] = None) -> None:
         """
         处理目录中的所有图片