chakcy/sophon_ai_project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

modify: 多线程管理器

Browse Source
This commit is contained in:
chakcy
2026-02-06 09:40:45 +08:00
parent 2ff9a2ecb1
commit 91d8456d4c
26 changed files with 81123 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
YOLOv8/postprocess_numpy.py
+1 -1
View File
@@ -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
+72 -64
View File
@@ -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}")
logs/yolo_test_0_20260206_091608.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_10_20260206_091613.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_11_20260206_091613.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_12_20260206_091614.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_13_20260206_091615.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_14_20260206_091616.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_15_20260206_091617.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_1_20260206_091608.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_2_20260206_091609.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_3_20260206_091609.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_4_20260206_091609.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_5_20260206_091610.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_6_20260206_091610.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_7_20260206_091611.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_8_20260206_091611.log
+5045
View File
File diff suppressed because it is too large Load Diff
logs/yolo_test_9_20260206_091612.log
+5045
View File
File diff suppressed because it is too large Load Diff
main.py
+23 -25
View File
@@ -1,33 +1,31 @@
import argparse
from sophon_chakcy.get_detector import get_detector
from sophon_chakcy.utils.detect_video_stream import DetectVideoStream
from threading import Thread
def detector_task(args):
detector_class, model_config = get_detector(args.name, True)
detector = detector_class(model_config, True, dev_id=args.dev_id, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh)
detect_video_stream = DetectVideoStream([detector], args.input)
t = Thread(target=detect_video_stream.run_task)
t.start()
import time
time.sleep(5)
detect_video_stream.stop_task()
t.join()
import time
from sophon_chakcy.utils.task_manager import TaskManager, TaskArgs
def main():
parser = argparse.ArgumentParser(description="YOLOv8目标检测")
parser.add_argument('--name', type=str, required=True, help='算法名称')
parser.add_argument('--input', type=str, required=True, help='输入图像或目录路径')
parser.add_argument('--dev_id', type=int, default=0, help='设备ID')
parser.add_argument('--conf_thresh', type=float, default=0.25, help='置信度阈值')
parser.add_argument('--nms_thresh', type=float, default=0.7, help='NMS阈值')
task_manager = TaskManager(16)
start_time = time.time()
print(start_time)
for i in range(16):
task_name = f"yolo_test_{i}"
args = TaskArgs(
task_name=task_name,
algorithms=["yolo_example", "yolo_example", "yolo_example"],
# input="rtsp://192.168.137.1:8554/stream",
input="./resources/datasets/test_car_person_1080P.mp4",
output=f"rtsp://192.168.137.1:8554/stream_{i}",
)
task_manager.submit(args)
args = parser.parse_args()
# time.sleep(30)
# for i in range(1):
# task_name = f"yolo_test_{i}"
# task_manager.stop_task(task_name)
# print(time.time() - start_time)
detector_task(args)
if __name__ == "__main__":
main()
main()
public/main.py
+20 -43
View File
@@ -1,57 +1,34 @@
import os
import argparse
from sophon_chakcy.get_detector import get_detector
import time
from sophon_chakcy.utils.task_manager import TaskManager, TaskArgs
from module_bank import PythonToSQLite
from multiprocessing import Process
# 安装导入器
packer = PythonToSQLite("plugins.mbank")
finder = packer.install_importer()
def detector_task(args):
detector_class, model_config = get_detector(args.name, True)
detector = detector_class(model_config, True, dev_id=args.dev_id, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh)
# 根据输入类型进行处理
if os.path.isfile(args.input):
# 单张图像
detector.detect_single_image(args.input)
elif os.path.isdir(args.input):
# 图像目录
detector.detect_images_in_directory(args.input)
else:
print(f"错误: {args.input} 既不是文件也不是目录")
def main():
parser = argparse.ArgumentParser(description="YOLOv8目标检测")
parser.add_argument('--name', type=str, required=True, help='算法名称')
parser.add_argument('--input', type=str, required=True, help='输入图像或目录路径')
parser.add_argument('--dev_id', type=int, default=0, help='设备ID')
parser.add_argument('--conf_thresh', type=float, default=0.25, help='置信度阈值')
parser.add_argument('--nms_thresh', type=float, default=0.7, help='NMS阈值')
task_manager = TaskManager(4)
start_time = time.time()
print(start_time)
for i in range(1):
task_name = f"yolo_test_{i}"
args = TaskArgs(
task_name=task_name,
algorithms=["yolo_example", "yolo_example"],
input="rtsp://192.168.137.1:8554/stream",
is_vfs=True
)
task_manager.submit(args)
args = parser.parse_args()
# time.sleep(5)
p = Process(target=detector_task,args=(args,))
p.start()
p.join()
# detector_class, model_config = get_detector(args.name, True)
# # 初始化检测器
# detector = detector_class(model_config, True, dev_id=args.dev_id, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh)
# # 根据输入类型进行处理
# if os.path.isfile(args.input):
# # 单张图像
# detector.detect_single_image(args.input)
# for i in range(4):
# task_name = f"yolo_test_{i}"
# task_manager.stop_task(task_name)
# print(time.time() - start_time)
# elif os.path.isdir(args.input):
# # 图像目录
# detector.detect_images_in_directory(args.input)
# else:
# print(f"错误: {args.input} 既不是文件也不是目录")
if __name__ == "__main__":
rtsp_test.log
+40
View File
@@ -0,0 +1,40 @@
BMvidDecCreateW5 board id 0 coreid 0
interleave TCP based rtp used: 0, 1
interleave TCP based rtp used: 2, 3
libbmvideo.so addr : /opt/sophon/libsophon-current/lib/libbmvideo.so.0, name_len: 34
vpu firmware addr: /opt/sophon/libsophon-current/lib/vpu_firmware/chagall_dec.bin
VERSION=0, REVISION=213135
ID[0], [ 1080 1920], RECON[ 0], FRM[ 46], FPS[46.00],[0.00] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 70], FPS[24.00],[0.00] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 94], FPS[24.00],[0.00] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 118], FPS[24.00],[33.16] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 141], FPS[23.00],[33.16] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 165], FPS[24.00],[33.16] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 188], FPS[23.00],[33.16] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 212], FPS[24.00],[27.54] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 236], FPS[24.00],[27.54] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 260], FPS[24.00],[27.54] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 284], FPS[24.00],[27.54] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 308], FPS[24.00],[26.27] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 332], FPS[24.00],[26.27] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 356], FPS[24.00],[26.27] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 380], FPS[24.00],[26.27] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 404], FPS[24.00],[25.68] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 428], FPS[24.00],[25.68] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 452], FPS[24.00],[25.68] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 476], FPS[24.00],[25.68] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 500], FPS[24.00],[25.31] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 524], FPS[24.00],[25.31] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 548], FPS[24.00],[25.31] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 572], FPS[24.00],[25.31] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 596], FPS[24.00],[25.31] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 620], FPS[24.00],[25.09] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 644], FPS[24.00],[25.09] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 668], FPS[24.00],[25.09] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 692], FPS[24.00],[25.09] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 716], FPS[24.00],[24.94] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 740], FPS[24.00],[24.94] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 764], FPS[24.00],[24.94] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 786], FPS[22.00],[24.94] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍ID[0], [ 1080 1920], RECON[ 0], FRM[ 811], FPS[25.00],[24.76] drop[ 0] | FRM[ 0], FPS[0.00],[0.00] enc_que_nums[0]
␍
Side by Side

After

Width:  |  Height:  |  Size: 5.0 KiB

src/sophon_chakcy/utils/detect_video_stream.py
+56 -48
View File
@@ -1,6 +1,9 @@
import sophon_chakcy.sail as sail
import threading
import time
from queue import Queue
from logging import getLogger, Logger
from typing import Optional
class DetectVideoStream:
@@ -11,7 +14,9 @@ class DetectVideoStream:
frame_skip: int = 0,
reconnect_attempts: int = 3,
reconnect_delay: int = 5,
is_save: bool = False):
is_save: bool = False,
output: str = "./result/output.mp4",
logger: Optional[Logger] = None):
"""
处理视频流（视频文件或RTSP流）
@@ -26,6 +31,10 @@ class DetectVideoStream:
Returns:
处理的帧数, 检测到的物体总数
"""
if logger is None:
self.logger = getLogger()
else:
self.logger = logger
self.detectors = detectors
self.dev_id = self.detectors[0].dev_id
self.handle = self.detectors[0].handle
@@ -48,7 +57,7 @@ class DetectVideoStream:
result_dir = r"results"
# 清理旧的目录
if os.path.exists(result_dir):
print(f"清理旧的构建目录: {result_dir}")
self.logger.info(f"清理旧的构建目录: {result_dir}")
shutil.rmtree(result_dir, ignore_errors=True)
# 创建目录
@@ -56,16 +65,14 @@ class DetectVideoStream:
if not result_dir_path.exists():
result_dir_path.mkdir(parents=True, exist_ok=True)
enc_params = f"width=1920:height=1080:bitrate=2000:gop=32:gop_preset=2:framerate=25"
self.encoder = sail.Encoder("results/output.mp4", self.dev_id, 'h264_bm', 'NV12', enc_params, 10)
def run_task(self):
if not self.decoder.is_opened():
print(f"错误: 无法打开视频流 {self.video_path}")
return 0, 0
enc_params = f"width=1920:height=1080:bitrate=2000:gop=32:gop_preset=2:framerate=12"
self.encoder = sail.Encoder(output, self.dev_id, 'h264_bm', 'NV12', enc_params, 10)
frame_shape = self.decoder.get_frame_shape()
print(f"视频帧形状: {frame_shape}")
def run_task(self) -> bool:
result_flog = True
if not self.decoder.is_opened():
self.logger.error(f"错误: 无法打开视频流 {self.video_path}")
result_flog = False
frame_count = 0
reconnect_count = 0
@@ -73,82 +80,83 @@ class DetectVideoStream:
try:
while not self.stop_event.is_set():
if self.max_frames > 0 and frame_count >= self.max_frames:
print(f"\n达到最大帧数限制: {self.max_frames}")
self.logger.warning(f"\n达到最大帧数限制: {self.max_frames}")
result_flog = True
break
# 读取一帧
bmimg = sail.BMImage()
ret = self.decoder.read(self.handle, bmimg)
if ret != 0:
print(f"读取帧失败，错误码: {ret}")
self.logger.error(f"读取帧失败，错误码: {ret}")
reconnect_count += 1
if reconnect_count > self.reconnect_attempts:
print("视频流无法继续读取")
self.logger.error("视频流无法继续读取")
break
print(f"尝试重连 ({reconnect_count}/{self.reconnect_attempts})...")
self.logger.info(f"尝试重连 ({reconnect_count}/{self.reconnect_attempts})...")
time.sleep(self.reconnect_delay)
self.decoder.release()
decoder = sail.Decoder(self.video_path, True, self.dev_id)
if decoder.is_opened():
reconnect_count = 0
print("重连成功，继续处理...")
self.logger.info("重连成功，继续处理...")
continue
else:
print("重连失败")
self.logger.warning("重连失败")
# time.sleep(0.01)
if self.frame_skip > 0 and frame_count % (self.frame_skip + 1) != 0:
frame_count += 1
if self.is_save:
self.encoder.video_write(bmimg)
continue
# 开始处理当前帧
start_time = time.time()
frame_count += 1
if frame_count % 10 == 0:
print(f"处理第 {frame_count} 帧...")
# 检测
for detector in self.detectors:
detections = detector.detect(bmimg)
self.logger.info(f"处理第 {frame_count} 帧...")
if len(detections) > 0:
detections = detections[0]
filtered_detections = []
for det in detections:
if len(det) >= 6:
x1, y1, x2, y2, conf, cls_id = det[:6]
if conf > detector.conf_thresh:
filtered_detections.append([x1, y1, x2, y2, conf, cls_id])
if len(filtered_detections) > 0:
result_img = detector.draw_detections(bmimg, filtered_detections, None)
else:
result_img = bmimg
else:
result_img = bmimg
filtered_detections_queue = Queue()
def thread_detect(bmimg, detector):
filtered_detections = detector.detect(bmimg)
filtered_detections_queue.put((detector, filtered_detections.copy()))
# 检测
threads = []
for detector in self.detectors:
t = threading.Thread(target=thread_detect, args=(bmimg, detector, ))
t.start()
threads.append(t)
for t in threads:
t.join(timeout=2)
while not filtered_detections_queue.empty():
detector, filtered_detections=filtered_detections_queue.get()
if len(filtered_detections) > 0:
bmimg = detector.draw_detections(bmimg, filtered_detections, None)
if self.is_save:
self.encoder.video_write(result_img)
self.encoder.video_write(bmimg) # type: ignore
# time.sleep(0.01)
except KeyboardInterrupt:
print("\n检测到键盘中断，停止处理")
self.logger.info("\n检测到键盘中断，停止处理")
result_flog = True
except Exception as e:
print(f"处理视频流时发生错误: {e}")
self.logger.error(f"处理视频流时发生错误: {e}")
result_flog = False
finally:
# 清理资源
self.decoder.release()
if self.is_save:
self.encoder.release()
print("视频解码器已释放")
self.logger.info("视频解码器已释放")
return result_flog
def stop_task(self):
print("提前关闭")
self.logger.info("提前关闭")
self.stop_event.set()
src/sophon_chakcy/utils/log_manager.py
+61
View File
@@ -0,0 +1,61 @@
import os
import time
import logging
import sys
from logging.handlers import RotatingFileHandler
class LogManager:
"""日志管理器"""
@staticmethod
def setup_logging(task_name=None):
"""配置日志系统"""
# 创建日志目录
log_dir = "logs"
if not os.path.exists(log_dir):
os.makedirs(log_dir)
# 生成日志文件名
if task_name:
log_file = os.path.join(log_dir, f"{task_name}_{time.strftime('%Y%m%d_%H%M%S')}.log")
logger_name = task_name
else:
log_file = os.path.join(log_dir, f"detection_{time.strftime('%Y%m%d_%H%M%S')}.log")
logger_name = "detection"
# 获取或创建特定名称的日志记录器
logger = logging.getLogger(logger_name)
logger.setLevel(logging.INFO)
# 清除现有处理器（避免重复添加）
if logger.handlers:
logger.handlers.clear()
# 文件处理器 - 轮转日志，最大10MB，保留5个备份
file_handler = RotatingFileHandler(
log_file,
maxBytes=10*1024*1024,
backupCount=5,
encoding='utf-8'
)
file_formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
file_handler.setFormatter(file_formatter)
logger.addHandler(file_handler)
# 控制台处理器 - 只输出警告及以上级别
console_handler = logging.StreamHandler(sys.stdout)
console_formatter = logging.Formatter(
'%(levelname)s - %(message)s'
)
console_handler.setFormatter(console_formatter)
console_handler.setLevel(logging.WARNING) # 只显示WARNING及以上
logger.addHandler(console_handler)
# 防止日志传递给父logger（如根logger）
logger.propagate = False
return log_file, logger # 同时返回logger实例
src/sophon_chakcy/utils/push_video_stream.py
+20
View File
@@ -0,0 +1,20 @@
import sophon_chakcy.sail as sail
from queue import Queue
class PushVideoStream:
def __init__(self, output, dev_id, width, height):
self.dev_id = dev_id
enc_params_dict = {
"width": width,
"height": height,
"bitrate": 2000,
"gop": 32,
"gop_preset": 2,
"framerate": 12
}
self.bmimg_queue = Queue(maxsize=30)
enc_params = ":".join({f"{key}:{value}" for key, value in enc_params_dict})
self.encoder = sail.Encoder(output, self.dev_id, 'h264_bm', 'NV12', enc_params, 10)
def
src/sophon_chakcy/utils/task_manager.py
+110
View File
@@ -0,0 +1,110 @@
from concurrent.futures import ThreadPoolExecutor
import threading
from dataclasses import dataclass, field
from .log_manager import LogManager
from .detect_video_stream import DetectVideoStream
from logging import getLogger
from ..get_detector import get_detector
from typing import List, Dict
@dataclass
class TaskArgs:
task_name: str = field(
metadata={"description": "任务名称"}
)
algorithms: list = field(
metadata={"description": "算法列表"}
)
input: str = field(
metadata={"description": "输入(一般为rtsp流地址或视频文件)"}
)
output: str = field(
metadata={"description": "输出(一般为rtsp流地址或视频文件)"}
)
is_vfs: bool = field(
default=False,
metadata={"description": "是否是虚拟文件系统(一般在生成环境中使用)"}
)
dev_id: int = field(
default=0,
metadata={"description": "设备ID"}
)
conf_thresh: float = field(
default=0.25,
metadata={"description": "置信度阈值"}
)
nms_thresh: float = field(
default=0.25,
metadata={"description": "NMS阈值"}
)
@classmethod
def from_dict(cls, data: dict) -> "TaskArgs":
data["algorithms"] = [item.strip() for item in data["algorithms"].split('*') if item.strip()]
print(data["algorithms"])
return cls(**data)
class TaskManager:
def __init__(self,
max_workers=None,
thread_name_prefix:str="task_"):
self.executor = ThreadPoolExecutor(
max_workers=max_workers,
thread_name_prefix=thread_name_prefix
)
self._task_dict: Dict[str, DetectVideoStream] = {}
self._results: List[bool] = []
self._lock = threading.Lock()
def submit(self, task_args: TaskArgs):
"""提交任务"""
if task_args.task_name in self._task_dict.keys():
print(f"'{task_args.task_name}' 任务已存在")
return
log_file, logger = LogManager.setup_logging(task_args.task_name)
logger.info(f"日志文件 {log_file}")
algorithms = [get_detector(algorithm, task_args.is_vfs) for algorithm in task_args.algorithms]
algorithms = [
algorithm[0](
algorithm[1],
task_args.is_vfs,
logger=logger,
dev_id=task_args.dev_id,
conf_thresh=task_args.conf_thresh,
nms_thresh=task_args.nms_thresh)
for
algorithm
in
algorithms
]
detect_video_stream = DetectVideoStream(algorithms,
task_args.input,
is_save=True,
logger=logger,
frame_skip=3,
output=task_args.output)
self._task_dict[task_args.task_name] = detect_video_stream
future = self.executor.submit(detect_video_stream.run_task)
future.add_done_callback(self._task_callback)
return future
def _task_callback(self, future):
"""任务完成回调"""
try:
result = future.result()
with self._lock:
self._results.append(result)
except Exception as e:
print(f"任务执行出错: {e}")
def stop_task(self, task_name):
self._task_dict[task_name].stop_task()
def shutdown(self, wait=True):
"""关闭线程池"""
self.executor.shutdown(wait=wait)
return self._results
src/sophon_chakcy/utils/task_manger.py
View File