modify: sophon-sail 创建管理系统
This commit is contained in:
@@ -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
|
||||
@@ -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}")
|
||||
Reference in New Issue
Block a user