sophon_ai_project/YOLOv8/yolov8_detector.py

import os
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


@dataclass
class YOLOv8Config:
    bmodel_path: str = field(
        metadata={"description": "模型路径"}
    )

    classes: list = field(
        metadata={"description": "检测类别"}
    )

    colors: list = field(
        metadata={"description": "边框颜色"}
    )

    @classmethod
    def from_dict(cls, data: dict) -> "YOLOv8Config":
        return cls(**data)


class Detector:
    """完整的YOLOv8检测器，包含可视化"""
    
    def __init__(self, 
                 handle:sail.Handle,
                 config_dict: dict, 
                 is_vfs=False, 
                 vfs_path: str="./application.svfs", 
                 conf_thresh = 0.25, 
                 nms_thresh = 0.7,
                 logger: Optional[Logger] = None):
        if logger is None:
            self.logger = getLogger()
        else:
            self.logger = logger
        self.handle = handle
        yolov8_config = YOLOv8Config.from_dict(config_dict)
        # 加载bmodel
        self.net = sail.Engine(self.handle)
        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"] # type: ignore
            self.net.load(bmodel_bytes, bmodel_size)
            # self.net = sail.Engine(bmodel_bytes, bmodel_size, dev_id, sail.IOMode.SYSO)
        else:
            self.net.load(yolov8_config.bmodel_path)
            # self.net = sail.Engine(yolov8_config.bmodel_path, dev_id, sail.IOMode.SYSO)
        graph_name = self.net.get_graph_names()[0]
        self.net.set_io_mode(graph_name, sail.IOMode.DEVIO)
        self.logger.info(f"加载模型: {yolov8_config.bmodel_path}")
        
        self.bmcv = sail.Bmcv(self.handle)
        self.graph_name = self.net.get_graph_names()[0]
        
        # 获取输入信息
        self.input_name = self.net.get_input_names(self.graph_name)[0]
        self.input_shape = self.net.get_input_shape(self.graph_name, self.input_name)
        self.input_dtype = self.net.get_input_dtype(self.graph_name, self.input_name)
        self.input_scale = self.net.get_input_scale(self.graph_name, self.input_name)
        
        # 获取img_dtype
        self.img_dtype = self.bmcv.get_bm_image_data_format(self.input_dtype)
        
        # 获取输出信息
        self.output_names = self.net.get_output_names(self.graph_name)
        self.output_tensors = {}
        for output_name in self.output_names:
            output_shape = self.net.get_output_shape(self.graph_name, output_name)
            output_dtype = self.net.get_output_dtype(self.graph_name, output_name)
            output = sail.Tensor(self.handle, output_shape, output_dtype, True, True)
            self.output_tensors[output_name] = output
        
        self.batch_size = self.input_shape[0]
        self.net_h = self.input_shape[2]
        self.net_w = self.input_shape[3]
        
        # 预处理参数
        self.ab = [x * self.input_scale / 255. for x in [1, 0, 1, 0, 1, 0]]
        
        # 后处理参数
        self.conf_thresh = conf_thresh
        self.nms_thresh = nms_thresh
        self.postprocess = PostProcess(
            conf_thresh=self.conf_thresh,
            nms_thresh=self.nms_thresh,
            agnostic=False,
            multi_label=False,
            max_det=300
        )
        
        self.colors = yolov8_config.colors
        self.classes = yolov8_config.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):
        """
        绘制检测结果到图像
        
        Args:
            image: sail.BMImage 对象
            detections: 检测结果列表，每个元素为 [x1, y1, x2, y2, confidence, class_id]
            save_path: 保存路径，如果不为None则保存图像
        
        Returns:
            绘制后的图像
        """
        if len(detections) == 0:
            self.logger.info("没有检测到任何物体")
            return image
        
        # 转换为BGR planar格式用于绘制
        img_bgr_planar = self.bmcv.convert_format(image)
        
        thickness = 2
        
        self.logger.info(f"\n绘制 {len(detections)} 个检测结果:")
        self.logger.info("-" * 80)
        
        for i, det in enumerate(detections):
            if len(det) >= 6:
                x1, y1, x2, y2, conf, cls_id = det[:6]
                
                # 转换为整数坐标
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                
                # 使用COLORS为不同类别分配颜色
                # 注意：class_id是0-79，但COLORS列表有160个颜色
                # 我们使用 (class_id + 1) % len(COLORS) 来确保不越界
                color_idx = (int(cls_id) + 1) % len(self.colors)
                color = self.colors[color_idx]
                
                # 获取类别名称
                class_name = self.classes[int(cls_id) + 1] if int(cls_id) + 1 < len(self.classes) else f"未知类别{int(cls_id)}"
                
                # 打印检测信息
                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,  # type: ignore
                                       (x2 - x1), (y2 - y1), 
                                       color, thickness)
        
        self.logger.info("-" * 80)
        
        # 保存图像
        if save_path:
            self.bmcv.imwrite(save_path, img_bgr_planar) # type: ignore
            self.logger.info(f"\n检测结果已保存到: {save_path}")
        
        return img_bgr_planar
    
    def preprocess(self, bmimg: sail.BMImage):
        """
        预处理图像

        Args:
            image: 输入图像 (HWC, BGR格式)

        Returns:
            (预处理后的BMImage, 原始尺寸, 缩放比例, 填充偏移)
        """
        # 转换为RGB
        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)
        
        # 调整大小并填充
        img_w = rgb_img.width()
        img_h = rgb_img.height()
        r = min(self.net_w / img_w, self.net_h / img_h)
        tw = int(round(r * img_w))
        th = int(round(r * img_h))
        tx1 = int(round((self.net_w - tw) / 2))
        ty1 = int(round((self.net_h - th) / 2))
        
        # 创建Padding属性
        attr = sail.PaddingAtrr()
        attr.set_stx(tx1)
        attr.set_sty(ty1)
        attr.set_w(tw)
        attr.set_h(th)
        attr.set_r(114)
        attr.set_g(114)
        attr.set_b(114)
        
        # 调整大小
        resized_img = self.bmcv.crop_and_resize_padding(rgb_img, 0, 0, img_w, img_h, 
                                                       self.net_w, self.net_h, attr)
        
        # 归一化
        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, 
                            ((self.ab[0], self.ab[1]), 
                             (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()
        self.logger.info(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)
        
        # 4. 执行推理
        input_tensors = {self.input_name: input_tensor}
        input_shapes = {self.input_name: self.input_shape}
        self.net.process(self.graph_name, input_tensors, input_shapes, self.output_tensors)
        
        # 5. 获取输出
        outputs = {}
        for name in self.output_names:
            outputs[name] = self.output_tensors[name].asnumpy()[0]
        
        # 6. 后处理
        output_key = list(outputs.keys())[0]
        pred = outputs[output_key]
        
        # 添加batch维度
        pred_with_batch = pred[np.newaxis, :, :]  # (1, 8400, 84)
        
        # 准备后处理参数
        org_size_list = [(orig_width, orig_height)]
        ratios_batch = [(r, r)]
        txy_batch = [(tx1, ty1)]
        
        # 执行后处理
        detections = self.postprocess([pred_with_batch], org_size_list, ratios_batch, txy_batch)
        
        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"):
        """
        检测单张图像
        
        Args:
            image_path: 图像路径
            output_dir: 输出目录
        
        Returns:
            检测结果列表
        """
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        self.logger.info(f"\n处理图像: {image_path}")
        
        # 1. 解码图像
        decoder = sail.Decoder(image_path, True)
        bmimg = sail.BMImage()
        ret = decoder.read(self.handle, bmimg)
        
        if ret != 0:
            self.logger.error(f"错误: 无法解码图像 {image_path}")
            return []
        
        orig_width = bmimg.width()
        orig_height = bmimg.height()
        self.logger.info(f"图像大小: {orig_width}x{orig_height}")
        
        filtered_detections = self.detect(bmimg)
        
        # 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)
            
            return filtered_detections
        else:
            self.logger.info("没有检测到任何物体")
            return []
    
    def detect_images_in_directory(self, input_dir, output_dir="./detection_results"):
        """
        检测目录中的所有图像
        
        Args:
            input_dir: 输入目录
            output_dir: 输出目录
        """
        if not os.path.exists(input_dir):
            self.logger.error(f"错误: 目录不存在 {input_dir}")
            return
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        self.logger.info(f"开始检测目录: {input_dir}")
        self.logger.info(f"结果保存到: {output_dir}")
        
        # 统计信息
        total_images = 0
        total_detections = 0
        
        # 遍历目录
        for root, dirs, files in os.walk(input_dir):
            for file in files:

                image_path = os.path.join(root, file)
                total_images += 1
                    
                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)
        
        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}")