modify: 多线程管理器

YOLOv8/postprocess_numpy.py

@@ -31,7 +31,7 @@ class PostProcess:
             # 1 output
             dets = np.concatenate(preds_batch)
         else:
-            print('preds_batch type: '.format(type(preds_batch)))
+            # print('preds_batch type: '.format(type(preds_batch)))
             raise NotImplementedError
 
         outs = self.nms.non_max_suppression(dets,

YOLOv8/yolov8_detector.py

@@ -1,13 +1,11 @@
-# yolov8_inference_with_colors.py
 import os
-import time
 import numpy as np
 import sophon_chakcy.sail as sail
 from .postprocess_numpy import PostProcess
 from sqlite_vfs.core import SQLiteVFS
 from dataclasses import dataclass, field
+from logging import Logger, getLogger
 from typing import Optional
-import threading
 
 
 @dataclass
@@ -32,18 +30,28 @@ class YOLOv8Config:
 class Detector:
     """完整的YOLOv8检测器，包含可视化"""
     
-    def __init__(self, config_dict: dict, is_vfs=False, vfs_path: str="./application.svfs", dev_id=0, conf_thresh = 0.25, nms_thresh = 0.7):
+    def __init__(self, 
+                 config_dict: dict, 
+                 is_vfs=False, 
+                 vfs_path: str="./application.svfs", 
+                 dev_id=0, conf_thresh = 0.25, 
+                 nms_thresh = 0.7,
+                 logger: Optional[Logger] = None):
+        if logger is None:
+            self.logger = getLogger()
+        else:
+            self.logger = logger
         self.dev_id = dev_id
         yolov8_config = YOLOv8Config.from_dict(config_dict)
         # 加载bmodel
         if is_vfs:
             vfs = SQLiteVFS(vfs_path)
             bmodel_bytes = vfs.read_file(yolov8_config.bmodel_path)
-            bmodel_size = vfs.get_file_info(yolov8_config.bmodel_path)["file_size"]
+            bmodel_size = vfs.get_file_info(yolov8_config.bmodel_path)["file_size"] # type: ignore
             self.net = sail.Engine(bmodel_bytes, bmodel_size, dev_id, sail.IOMode.SYSO)
         else:
             self.net = sail.Engine(yolov8_config.bmodel_path, dev_id, sail.IOMode.SYSO)
-        print(f"加载模型: {yolov8_config.bmodel_path}")
+        self.logger.info(f"加载模型: {yolov8_config.bmodel_path}")
         
         self.handle = sail.Handle(dev_id)
         self.bmcv = sail.Bmcv(self.handle)
@@ -88,11 +96,11 @@ class Detector:
         self.colors = yolov8_config.colors
         self.classes = yolov8_config.classes
         
-        print(f"模型信息:")
-        print(f"  输入形状: {self.input_shape}")
-        print(f"  输入类型: {self.input_dtype}")
-        print(f"  colors数量: {len(self.colors)}")
-        print(f"  classes类别数: {len(self.classes)}")
+        self.logger.info(f"模型信息:")
+        self.logger.info(f"  输入形状: {self.input_shape}")
+        self.logger.info(f"  输入类型: {self.input_dtype}")
+        self.logger.info(f"  colors数量: {len(self.colors)}")
+        self.logger.info(f"  classes类别数: {len(self.classes)}")
     
     def draw_detections(self, image, detections, save_path=None):
         """
@@ -107,7 +115,7 @@ class Detector:
             绘制后的图像
         """
         if len(detections) == 0:
-            print("没有检测到任何物体")
+            self.logger.info("没有检测到任何物体")
             return image
         
         # 转换为BGR planar格式用于绘制
@@ -115,8 +123,8 @@ class Detector:
         
         thickness = 2
         
-        print(f"\n绘制 {len(detections)} 个检测结果:")
-        print("-" * 80)
+        self.logger.info(f"\n绘制 {len(detections)} 个检测结果:")
+        self.logger.info("-" * 80)
         
         for i, det in enumerate(detections):
             if len(det) >= 6:
@@ -135,20 +143,20 @@ class Detector:
                 class_name = self.classes[int(cls_id) + 1] if int(cls_id) + 1 < len(self.classes) else f"未知类别{int(cls_id)}"
                 
                 # 打印检测信息
-                print(f"检测 {i+1:2d}: {class_name:15s} | 置信度: {conf:.4f} | 位置: [{x1:4d}, {y1:4d}, {x2:4d}, {y2:4d}] | 大小: {x2-x1:4d}x{y2-y1:4d} | 颜色: {color}")
+                self.logger.info(f"检测 {i+1:2d}: {class_name:15s} | 置信度: {conf:.4f} | 位置: [{x1:4d}, {y1:4d}, {x2:4d}, {y2:4d}] | 大小: {x2-x1:4d}x{y2-y1:4d} | 颜色: {color}")
                 
                 # 绘制边界框
                 if (x2 - x1) > thickness * 2 and (y2 - y1) > thickness * 2:
-                    self.bmcv.rectangle(img_bgr_planar, x1, y1, 
+                    self.bmcv.rectangle(img_bgr_planar, x1, y1,  # type: ignore
                                        (x2 - x1), (y2 - y1), 
                                        color, thickness)
         
-        print("-" * 80)
+        self.logger.info("-" * 80)
         
         # 保存图像
         if save_path:
-            self.bmcv.imwrite(save_path, img_bgr_planar)
-            print(f"\n检测结果已保存到: {save_path}")
+            self.bmcv.imwrite(save_path, img_bgr_planar) # type: ignore
+            self.logger.info(f"\n检测结果已保存到: {save_path}")
         
         return img_bgr_planar
     
@@ -163,8 +171,8 @@ class Detector:
             (预处理后的BMImage, 原始尺寸, 缩放比例, 填充偏移)
         """
         # 转换为RGB
-        rgb_img = sail.BMImage(self.handle, bmimg.height(), bmimg.width(),
-                               sail.Format.FORMAT_RGB_PLANAR, sail.DATA_TYPE_EXT_1N_BYTE)
+        rgb_img = sail.BMImage(self.handle, bmimg.height(), bmimg.width(), # type: ignore
+                               sail.Format.FORMAT_RGB_PLANAR, sail.DATA_TYPE_EXT_1N_BYTE) # type: ignore
         self.bmcv.convert_format(bmimg, rgb_img)
         
         # 调整大小并填充
@@ -191,8 +199,8 @@ class Detector:
                                                        self.net_w, self.net_h, attr)
         
         # 归一化
-        preprocessed_img = sail.BMImage(self.handle, self.net_h, self.net_w,
-                                        sail.Format.FORMAT_RGB_PLANAR, self.img_dtype)
+        preprocessed_img = sail.BMImage(self.handle, self.net_h, self.net_w, # type: ignore
+                                        sail.Format.FORMAT_RGB_PLANAR, self.img_dtype) # type: ignore
         
         # 转换并归一化
         self.bmcv.convert_to(resized_img, preprocessed_img, 
@@ -205,7 +213,7 @@ class Detector:
     def detect(self, bmimg: sail.BMImage):
         orig_width = bmimg.width()
         orig_height = bmimg.height()
-        print(f"图像大小: {orig_width}x{orig_height}")
+        self.logger.info(f"图像大小: {orig_width}x{orig_height}")
         
         # 2. 预处理
         preprocessed_img, r, tx1, ty1 = self.preprocess(bmimg)
@@ -238,7 +246,23 @@ class Detector:
         
         # 执行后处理
         detections = self.postprocess([pred_with_batch], org_size_list, ratios_batch, txy_batch)
-        return detections
+        
+        filtered_detections = []
+
+        if len(detections) > 0:
+            detections = detections[0]  # 取第一个batch的结果
+            
+            # 7. 过滤低置信度的检测
+            for det in detections:
+                if len(det) >= 6:
+                    x1, y1, x2, y2, conf, cls_id = det[:6]
+                    if conf > self.conf_thresh:
+                        filtered_detections.append([x1, y1, x2, y2, conf, cls_id])
+            
+            self.logger.info(f"原始检测数量: {len(detections)}")
+            self.logger.info(f"过滤后数量 (置信度 > {self.conf_thresh}): {len(filtered_detections)}")
+
+        return filtered_detections
 
     def detect_single_image(self, image_path, output_dir="./detection_results"):
         """
@@ -254,7 +278,7 @@ class Detector:
         # 创建输出目录
         os.makedirs(output_dir, exist_ok=True)
         
-        print(f"\n处理图像: {image_path}")
+        self.logger.info(f"\n处理图像: {image_path}")
         
         # 1. 解码图像
         decoder = sail.Decoder(image_path, True, self.dev_id)
@@ -262,44 +286,28 @@ class Detector:
         ret = decoder.read(self.handle, bmimg)
         
         if ret != 0:
-            print(f"错误: 无法解码图像 {image_path}")
+            self.logger.error(f"错误: 无法解码图像 {image_path}")
             return []
         
         orig_width = bmimg.width()
         orig_height = bmimg.height()
-        print(f"图像大小: {orig_width}x{orig_height}")
+        self.logger.info(f"图像大小: {orig_width}x{orig_height}")
         
-        detections = self.detect(bmimg)
+        filtered_detections = self.detect(bmimg)
         
-        if len(detections) > 0:
-            detections = detections[0]  # 取第一个batch的结果
-            
-            # 7. 过滤低置信度的检测
-            conf_threshold = 0.25
-            filtered_detections = []
-            
-            for det in detections:
-                if len(det) >= 6:
-                    x1, y1, x2, y2, conf, cls_id = det[:6]
-                    if conf > conf_threshold:
-                        filtered_detections.append([x1, y1, x2, y2, conf, cls_id])
-            
-            print(f"原始检测数量: {len(detections)}")
-            print(f"过滤后数量 (置信度 > {conf_threshold}): {len(filtered_detections)}")
-            
-            # 8. 绘制检测结果
-            if len(filtered_detections) > 0:
-                # 准备输出文件名
-                base_name = os.path.basename(image_path)
-                name_without_ext = os.path.splitext(base_name)[0]
-                output_path = os.path.join(output_dir, f"detected_{base_name}")
+        # 8. 绘制检测结果
+        if len(filtered_detections) > 0:
+            # 准备输出文件名
+            base_name = os.path.basename(image_path)
+            # name_without_ext = os.path.splitext(base_name)[0]
+            output_path = os.path.join(output_dir, f"detected_{base_name}")
                 
-                # 绘制检测框
-                self.draw_detections(bmimg, filtered_detections, output_path)
+            # 绘制检测框
+            self.draw_detections(bmimg, filtered_detections, output_path)
             
             return filtered_detections
         else:
-            print("没有检测到任何物体")
+            self.logger.info("没有检测到任何物体")
             return []
     
     def detect_images_in_directory(self, input_dir, output_dir="./detection_results"):
@@ -311,14 +319,14 @@ class Detector:
             output_dir: 输出目录
         """
         if not os.path.exists(input_dir):
-            print(f"错误: 目录不存在 {input_dir}")
+            self.logger.error(f"错误: 目录不存在 {input_dir}")
             return
         
         # 创建输出目录
         os.makedirs(output_dir, exist_ok=True)
         
-        print(f"开始检测目录: {input_dir}")
-        print(f"结果保存到: {output_dir}")
+        self.logger.info(f"开始检测目录: {input_dir}")
+        self.logger.info(f"结果保存到: {output_dir}")
         
         # 统计信息
         total_images = 0
@@ -331,15 +339,15 @@ class Detector:
                 image_path = os.path.join(root, file)
                 total_images += 1
                     
-                print(f"\n{'='*80}")
-                print(f"处理图像 {total_images}: {file}")
+                self.logger.info(f"\n{'='*80}")
+                self.logger.info(f"处理图像 {total_images}: {file}")
                     
                 # 检测单张图像
                 detections = self.detect_single_image(image_path, output_dir)
                 total_detections += len(detections)
         
-        print(f"\n{'='*80}")
-        print(f"检测完成!")
-        print(f"总处理图像数: {total_images}")
-        print(f"总检测物体数: {total_detections}")
-        print(f"平均每张图像检测数: {total_detections/max(total_images, 1):.2f}")
+        self.logger.info(f"\n{'='*80}")
+        self.logger.info(f"检测完成!")
+        self.logger.info(f"总处理图像数: {total_images}")
+        self.logger.info(f"总检测物体数: {total_detections}")
+        self.logger.info(f"平均每张图像检测数: {total_detections/max(total_images, 1):.2f}")