modify: sophon-sail 创建管理系统

This commit is contained in:
2026-02-03 21:32:23 +08:00
parent b18bea8c5b
commit 0b32a1f70b
27 changed files with 4743 additions and 1132 deletions
View File
+190
View File
@@ -0,0 +1,190 @@
#===----------------------------------------------------------------------===#
#
# Copyright (C) 2022 Sophgo Technologies Inc. All rights reserved.
#
# SOPHON-DEMO is licensed under the 2-Clause BSD License except for the
# third-party components.
#
#===----------------------------------------------------------------------===#
import numpy as np
class PostProcess:
def __init__(self, conf_thresh=0.001, nms_thresh=0.7, agnostic=False, multi_label=True, max_det=300):
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.agnostic_nms = agnostic
self.multi_label = multi_label
self.max_det = max_det
self.nms = pseudo_torch_nms()
def __call__(self, preds_batch, org_size_batch, ratios_batch, txy_batch):
"""
post-processing
:param preds_batch: list of predictions in a batch
:param org_size_batch: list of (org_img_w, org_img_h) in a batch
:param ratios_batch: list of (ratio_x, ratio_y) in a batch when resize-and-center-padding
:param txy_batch: list of (tx, ty) in a batch when resize-and-center-padding
:return:
"""
if isinstance(preds_batch, list) and len(preds_batch) == 1:
# 1 output
dets = np.concatenate(preds_batch)
else:
print('preds_batch type: '.format(type(preds_batch)))
raise NotImplementedError
outs = self.nms.non_max_suppression(dets,
self.conf_thresh,
self.nms_thresh,
agnostic=False,
max_det=300,
multi_label=self.multi_label,
classes=None)
# Rescale boxes from img_size to im0 size
for det, (org_w, org_h), ratio, (tx1, ty1) in zip(outs, org_size_batch, ratios_batch, txy_batch):
if len(det):
# Rescale boxes from img_size to im0 size
coords = det[:, :4]
coords[:, [0, 2]] -= tx1 # x padding
coords[:, [1, 3]] -= ty1 # y padding
coords[:, [0, 2]] /= ratio[0]
coords[:, [1, 3]] /= ratio[1]
coords[:, [0, 2]] = coords[:, [0, 2]].clip(0, org_w - 1) # x1, x2
coords[:, [1, 3]] = coords[:, [1, 3]].clip(0, org_h - 1) # y1, y2
det[:, :4] = coords
return outs
# numpy multiclass nms implementation from original yolov5 repo torch implementation
class pseudo_torch_nms:
def nms_boxes(self, boxes, scores, iou_thres):
x = boxes[:, 0]
y = boxes[:, 1]
w = boxes[:, 2] - boxes[:, 0]
h = boxes[:, 3] - boxes[:, 1]
areas = w * h
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x[i], x[order[1:]])
yy1 = np.maximum(y[i], y[order[1:]])
xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])
w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
inter = w1 * h1
ovr = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(ovr <= iou_thres)[0]
order = order[inds + 1] # type: ignore
keep = np.array(keep)
return keep
def xywh2xyxy(self, x):
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
y = x.copy() if isinstance(x, np.ndarray) else np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
return y
def nms(self, pred, conf_thres=0.001, iou_thres=0.5, agnostic=False, max_det=1000):
return self.non_max_suppression(pred, conf_thres, iou_thres,
classes=None,
agnostic=agnostic,
multi_label=True,
max_det=max_det)
def non_max_suppression(self,
prediction,
conf_thres=0.25,
iou_thres=0.5,
classes=None,
agnostic=False,
multi_label=False,
labels=(),
max_det=300,
nm=0):
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping bounding boxes
Returns:
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
"""
bs = prediction.shape[0] # batch size
nc = prediction.shape[2] - nm - 4 # number of classes
mi = 4 + nc # mask start index
xc = prediction[:, :, 4:mi].max(2) > conf_thres # candidates
# Checks
# assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
# assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
# Settings
max_wh = 7680 # (pixels) maximum box width and height
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
# output = [torch.zeros((0, 6), device=prediction.device)] * bs
output = [np.zeros((0, 6 + nm))] * bs
for xi, x in enumerate(prediction): # image index, image inference
# Apply constraints
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
x = x[xc[xi]] # confidence
# If none remain process next image
if not x.shape[0]:
continue
# Detections matrix nx6 (xyxy, conf, cls)
box, cls, mask = x[:,:4], x[:, 4:nc+4], x[:, nc+4:]
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
box = self.xywh2xyxy(box)
# Detections matrix nx6 (xyxy, conf, cls)
if multi_label:
i, j = (cls > conf_thres).nonzero()
x = np.concatenate((box[i], x[i, 4 + j, None], j[:, None].astype(np.float32), mask[i]), 1)
else: # best class only
conf = cls.max(1, keepdims=True)
j_argmax = cls.argmax(1)
j = j_argmax if j_argmax.shape == x[:, 5:].shape else \
np.expand_dims(j_argmax, 1) # for argmax(axis, keepdims=True)
x = np.concatenate((box, conf, j.astype(np.float32), mask), 1)[conf.reshape(-1) > conf_thres]
# Check shape
n = x.shape[0] # number of boxes
if not n: # no boxes
continue
# elif n > max_nms: # excess boxes
# x_argsort = np.argsort(x[:, 4])[:max_nms] # sort by confidence
# x = x[x_argsort]
x_argsort = np.argsort(x[:, 4])[::-1][:max_nms] # sort by confidence
x = x[x_argsort]
# Batched NMS
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
#############################################
# i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
i = self.nms_boxes(boxes, scores, iou_thres)
############################################
if i.shape[0] > max_det: # limit detections
i = i[:max_det]
output[xi] = x[i]
return output
+320
View File
@@ -0,0 +1,320 @@
# yolov8_inference_with_colors.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
@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, config_dict: dict, is_vfs=False, vfs_path: str="./application.svfs", dev_id=0, conf_thresh = 0.25, nms_thresh = 0.7):
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"]
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.handle = sail.Handle(dev_id)
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
print(f"模型信息:")
print(f" 输入形状: {self.input_shape}")
print(f" 输入类型: {self.input_dtype}")
print(f" colors数量: {len(self.colors)}")
print(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:
print("没有检测到任何物体")
return image
# 转换为BGR planar格式用于绘制
img_bgr_planar = self.bmcv.convert_format(image)
thickness = 2
print(f"\n绘制 {len(detections)} 个检测结果:")
print("-" * 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)}"
# 打印检测信息
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}")
# 绘制边界框
if (x2 - x1) > thickness * 2 and (y2 - y1) > thickness * 2:
self.bmcv.rectangle(img_bgr_planar, x1, y1,
(x2 - x1), (y2 - y1),
color, thickness)
print("-" * 80)
# 保存图像
if save_path:
self.bmcv.imwrite(save_path, img_bgr_planar)
print(f"\n检测结果已保存到: {save_path}")
return img_bgr_planar
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}")
# 2. 预处理
# 转换为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)
# 调整大小并填充
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,
sail.Format.FORMAT_RGB_PLANAR, self.img_dtype)
# 转换并归一化
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])))
# 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)
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}")
# 绘制检测框
self.draw_detections(bmimg, filtered_detections, output_path)
return filtered_detections
else:
print("没有检测到任何物体")
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):
print(f"错误: 目录不存在 {input_dir}")
return
# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
print(f"开始检测目录: {input_dir}")
print(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
print(f"\n{'='*80}")
print(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}")