modify: 实现视频流推理

YOLOv8/yolov8_detector.py

@@ -1,11 +1,14 @@
 # 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 typing import Optional
+import threading
+
 
 @dataclass
 class YOLOv8Config:
@@ -149,37 +152,17 @@ class Detector:
         
         return img_bgr_planar
     
-    def detect_single_image(self, image_path, output_dir="./detection_results"):
+    def preprocess(self, bmimg: sail.BMImage):
         """
-        检测单张图像
-        
+        预处理图像
+
         Args:
-            image_path: 图像路径
-            output_dir: 输出目录
-        
+            image: 输入图像 (HWC, BGR格式)
+
         Returns:
-            检测结果列表
+            (预处理后的BMImage, 原始尺寸, 缩放比例, 填充偏移)
         """
-        # 创建输出目录
-        os.makedirs(output_dir, exist_ok=True)
-        
-        print(f"\n处理图像: {image_path}")
-        
-        # 1. 解码图像
-        decoder = sail.Decoder(image_path, True, self.dev_id)
-        bmimg = sail.BMImage()
-        ret = decoder.read(self.handle, bmimg)
-        
-        if ret != 0:
-            print(f"错误: 无法解码图像 {image_path}")
-            return []
-        
-        orig_width = bmimg.width()
-        orig_height = bmimg.height()
-        print(f"图像大小: {orig_width}x{orig_height}")
-        
-        # 2. 预处理
-        # 转换为RGB格式
+        # 转换为RGB
         rgb_img = sail.BMImage(self.handle, bmimg.height(), bmimg.width(),
                                sail.Format.FORMAT_RGB_PLANAR, sail.DATA_TYPE_EXT_1N_BYTE)
         self.bmcv.convert_format(bmimg, rgb_img)
@@ -217,6 +200,16 @@ class Detector:
                              (self.ab[2], self.ab[3]), 
                              (self.ab[4], self.ab[5])))
         
+        return preprocessed_img, r, tx1, ty1
+
+    def detect(self, bmimg: sail.BMImage):
+        orig_width = bmimg.width()
+        orig_height = bmimg.height()
+        print(f"图像大小: {orig_width}x{orig_height}")
+        
+        # 2. 预处理
+        preprocessed_img, r, tx1, ty1 = self.preprocess(bmimg)
+
         # 3. 准备输入张量
         input_tensor = sail.Tensor(self.handle, self.input_shape, self.input_dtype, False, False)
         self.bmcv.bm_image_to_tensor(preprocessed_img, input_tensor)
@@ -245,6 +238,38 @@ class Detector:
         
         # 执行后处理
         detections = self.postprocess([pred_with_batch], org_size_list, ratios_batch, txy_batch)
+        return detections
+
+    def detect_single_image(self, image_path, output_dir="./detection_results"):
+        """
+        检测单张图像
+        
+        Args:
+            image_path: 图像路径
+            output_dir: 输出目录
+        
+        Returns:
+            检测结果列表
+        """
+        # 创建输出目录
+        os.makedirs(output_dir, exist_ok=True)
+        
+        print(f"\n处理图像: {image_path}")
+        
+        # 1. 解码图像
+        decoder = sail.Decoder(image_path, True, self.dev_id)
+        bmimg = sail.BMImage()
+        ret = decoder.read(self.handle, bmimg)
+        
+        if ret != 0:
+            print(f"错误: 无法解码图像 {image_path}")
+            return []
+        
+        orig_width = bmimg.width()
+        orig_height = bmimg.height()
+        print(f"图像大小: {orig_width}x{orig_height}")
+        
+        detections = self.detect(bmimg)
         
         if len(detections) > 0:
             detections = detections[0]  # 取第一个batch的结果