From 540a20e15bb200ac1967eb65c498915a2fe05bdc Mon Sep 17 00:00:00 2001 From: yooooger <761181201@qq.com> Date: Mon, 9 Mar 2026 14:22:19 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E4=B8=BAyolo=E5=AE=98?= =?UTF-8?q?=E6=96=B9=E6=8F=90=E4=BE=9B=E7=9A=84=E8=BE=93=E5=87=BA=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E6=A0=BC=E5=BC=8F=EF=BC=8C=E4=BF=9D=E5=AD=98=E6=A0=87?= =?UTF-8?q?=E7=AD=BE=E6=96=87=E4=BB=B6=E6=97=B6=E7=9B=B4=E6=8E=A5=E4=BB=8E?= =?UTF-8?q?yolo=E7=9A=84=E4=B8=B4=E6=97=B6=E8=BE=93=E5=87=BA=E7=9B=AE?= =?UTF-8?q?=E5=BD=95=E5=A4=8D=E5=88=B6txt=E6=96=87=E4=BB=B6=EF=BC=8C?= =?UTF-8?q?=E8=80=8C=E4=B8=8D=E6=98=AF=E9=87=8D=E6=96=B0=E7=94=9F=E6=88=90?= =?UTF-8?q?=E6=A0=87=E7=AD=BE=E6=96=87=E4=BB=B6=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- predict/predict_yolo11seg.py | 65 +++++++++++++----------------------- 1 file changed, 23 insertions(+), 42 deletions(-) diff --git a/predict/predict_yolo11seg.py b/predict/predict_yolo11seg.py index da34098..99746c7 100644 --- a/predict/predict_yolo11seg.py +++ b/predict/predict_yolo11seg.py @@ -91,6 +91,7 @@ class InferenceResult: self.scores = [] # 置信度列表 self.class_names = [] # 类别名称列表 self.inference_time = 0.0 # 推理时间（秒） + self.temp_label_path = None # YOLO临时输出目录路径 class YOLOSegmentationInference: @@ -290,8 +291,12 @@ class YOLOSegmentationInference: print(f"正在处理图片: {os.path.basename(image_path)}") start_time = time.time() - # 使用YOLO模型进行推理 - predictions = self.model(image_path, conf=conf_threshold, iou=iou_threshold)[0] + # 使用YOLO模型进行推理（启用save_txt自动生成标签） + temp_output_dir = os.path.join(os.path.dirname(image_path), ".yolo_temp") + predictions = self.model(image_path, conf=conf_threshold, iou=iou_threshold, + save=False, save_txt=True, project=temp_output_dir, + name="labels", exist_ok=True)[0] + result.temp_label_path = temp_output_dir result.inference_time = time.time() - start_time @@ -393,7 +398,7 @@ class YOLOSegmentationInference: result: 推理结果 output_dir: 输出目录 save_mask: 是否保存单独的掩码文件 - save_label: 是否保存YOLO格式的标签文件 + save_label: 是否保存YOLO格式的标签文件（使用YOLO原生格式） """ if result.result_image is None: return @@ -432,50 +437,26 @@ class YOLOSegmentationInference: print(f"共保存 {len(result.masks)} 个掩码文件到: {mask_dir}") - # 保存YOLO格式的标签文件 - if save_label and result.masks is not None and len(result.masks) > 0 and len(result.boxes) > 0: - # label_dir = os.path.join(output_dir, "labels") + # 保存YOLO格式的标签文件（直接从YOLO输出复制） + if save_label and result.masks is not None and len(result.masks) > 0: label_dir = output_dir os.makedirs(label_dir, exist_ok=True) label_path = os.path.join(label_dir, f"{base_name}.txt") - - with open(label_path, 'w') as f: - for i in range(len(result.masks)): - class_id = result.classes[i] - score = result.scores[i] - mask = result.masks[i] - - # 获取掩码的多边形轮廓 - contours, _ = cv2.findContours((mask > 0.5).astype(np.uint8), - cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - - if contours: - # 取最大的轮廓 - contour = max(contours, key=cv2.contourArea) - - # 简化轮廓 - epsilon = 0.001 * cv2.arcLength(contour, True) - approx = cv2.approxPolyDP(contour, epsilon, True) - - # 归一化坐标 - h, w = mask.shape - points = [] - for point in approx: - x = max(0.0, min(1.0, point[0][0] / w)) - y = max(0.0, min(1.0, point[0][1] / h)) - points.extend([x, y]) - - # 写入标签文件 - if len(points) >= 6: # 至少3个点 - line = f"{class_id} {' '.join(map(lambda x: f'{x:.6f}', points))}\n" - f.write(line) - - print(f"标签文件已保存: {label_path}") - redir_obj["label"] = label_path + + # 从YOLO临时输出目录复制txt文件 + if hasattr(result, 'temp_label_path') and result.temp_label_path: + yolo_label_path = os.path.join(result.temp_label_path, "labels", f"{base_name}.txt") + if os.path.exists(yolo_label_path): + shutil.copy(yolo_label_path, label_path) + print(f"标签文件已保存: {label_path}") + redir_obj["label"] = label_path + else: + print(f"警告: YOLO生成的标签文件不存在: {yolo_label_path}") + else: + print(f"警告: 未找到YOLO临时输出目录") result_save.append(redir_obj) - - + except PermissionError: print(f"权限错误: 无法写入到目录 {output_dir}") except Exception as e: