yipai-tech/SupervisorAI
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

print改为logger打印

Browse Source
This commit is contained in:
zqc
2026-02-03 09:15:37 +08:00
parent 42d7d24c4d
commit f6ea6082f3
5 changed files with 46 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
algorithm/checkpoint/npu_yolo_onnx_person_car_phone.py
View File

@@ -5,6 +5,9 @@ import onnxruntime as ort
import os import os
import time import time
from utils.logger import get_logger
logger = get_logger(__name__)
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114)): def letterbox(img, new_shape=(640, 640), color=(114, 114, 114)):
shape = img.shape[:2] # h, w shape = img.shape[:2] # h, w
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
@@ -35,23 +38,23 @@ class YOLOv8_ONNX:
self.session = ort.InferenceSession(onnx_path, providers=providers) self.session = ort.InferenceSession(onnx_path, providers=providers)
actual_providers = self.session.get_providers() actual_providers = self.session.get_providers()
print("YOLO Providers:", actual_providers) logger.info("YOLO Providers:", actual_providers)
if "CANNExecutionProvider" in actual_providers: if "CANNExecutionProvider" in actual_providers:
print("[INFO] YOLO 使用 CANNExecutionProvider（昇腾 NPU）") logger.info("[INFO] YOLO 使用 CANNExecutionProvider（昇腾 NPU）")
elif 'CUDAExecutionProvider' in actual_providers: elif 'CUDAExecutionProvider' in actual_providers:
print("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）") logger.info("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
else: else:
print("[INFO] YOLO 使用 CPUExecutionProvider") logger.info("[INFO] YOLO 使用 CPUExecutionProvider")
self.conf_threshold = conf_threshold self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold self.iou_threshold = iou_threshold
self.input_name = self.session.get_inputs()[0].name self.input_name = self.session.get_inputs()[0].name
self.input_size = (input_size, input_size) if isinstance(input_size, int) else input_size self.input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
print(f"模型输入名称: {self.input_name}") logger.info(f"模型输入名称: {self.input_name}")
print(f"模型输入形状: {self.session.get_inputs()[0].shape}") logger.info(f"模型输入形状: {self.session.get_inputs()[0].shape}")
print(f"模型输出形状: {self.session.get_outputs()[0].shape}") logger.info(f"模型输出形状: {self.session.get_outputs()[0].shape}")
def preprocess(self, img): def preprocess(self, img):
self.orig_shape = img.shape[:2] self.orig_shape = img.shape[:2]

17
algorithm/checkpoint/npu_yolo_pose_onnx.py
View File

@@ -9,7 +9,8 @@ import cv2
import numpy as np import numpy as np
import onnxruntime as ort import onnxruntime as ort
from utils.logger import get_logger
logger = get_logger(__name__)
# ------------------------------------------------- # -------------------------------------------------
@@ -89,14 +90,14 @@ class YOLOv8_Pose_ONNX:
# 获取真实工作 provider # 获取真实工作 provider
actual_providers = self.session.get_providers() actual_providers = self.session.get_providers()
print("YOLO Providers:", actual_providers) logger.info("YOLO Providers:", actual_providers)
if "CANNExecutionProvider" in actual_providers: if "CANNExecutionProvider" in actual_providers:
print("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）") logger.info("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）")
elif 'CUDAExecutionProvider' in actual_providers: elif 'CUDAExecutionProvider' in actual_providers:
print("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）") logger.info("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
else: else:
print("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）") logger.info("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）")
self.conf_threshold = conf_threshold self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold self.iou_threshold = iou_threshold
@@ -104,9 +105,9 @@ class YOLOv8_Pose_ONNX:
self.input_name = self.session.get_inputs()[0].name self.input_name = self.session.get_inputs()[0].name
self.input_size = (input_size, input_size) self.input_size = (input_size, input_size)
print(f"模型输入名称: {self.input_name}") logger.info(f"模型输入名称: {self.input_name}")
print(f"模型输入形状: {self.session.get_inputs()[0].shape}") logger.info(f"模型输入形状: {self.session.get_inputs()[0].shape}")
print(f"模型输出形状: {self.session.get_outputs()[0].shape}") logger.info(f"模型输出形状: {self.session.get_outputs()[0].shape}")
def nms(self, boxes, scores, iou_threshold=0.45): def nms(self, boxes, scores, iou_threshold=0.45):

33
rtsp_service_ws_kadian.py
View File

@@ -16,6 +16,9 @@ from biz.checkpoint.checkpoint_biz import KadianDetector, RTSP_TARGET_FPS, ALERT
from test_cam import get_camera_preview_url from test_cam import get_camera_preview_url
from utils.web_socket_sender import WebSocketSender from utils.web_socket_sender import WebSocketSender
from utils.logger import get_logger
logger = get_logger(__name__)
WS_HOST = "0.0.0.0" WS_HOST = "0.0.0.0"
WS_PORT = 8765 WS_PORT = 8765
@@ -49,22 +52,22 @@ class RTSPCaptureWorker(threading.Thread):
try: try:
if self.reconnect_count >= self.max_reconnects: if self.reconnect_count >= self.max_reconnects:
print(f"[WARN] RTSP: {self.camera_cfg.name} reach max reconnects, refresh url") logger.warning(f"[WARN] RTSP: {self.camera_cfg.name} reach max reconnects, refresh url")
self.reconnect_count = 0 self.reconnect_count = 0
new_url = self.refresh_video_url() new_url = self.refresh_video_url()
if new_url: if new_url:
self.rtsp_url = new_url self.rtsp_url = new_url
else: else:
print(f"[ERROR] refresh RTSP URL is empty, do nothing") logger.error(f"[ERROR] refresh RTSP URL is empty, do nothing")
# 检查rtsp_url是否为空或None，如果是则重新获取 # 检查rtsp_url是否为空或None，如果是则重新获取
if not self.rtsp_url: if not self.rtsp_url:
print(f"[WARN] RTSP URL is empty, refreshing...") logger.warning(f"[WARN] RTSP URL is empty, refreshing...")
new_url = self.refresh_video_url() new_url = self.refresh_video_url()
if new_url: if new_url:
self.rtsp_url = new_url self.rtsp_url = new_url
else: else:
print(f"[ERROR] RTSP URL is still empty, retrying in 5 seconds") logger.error(f"[ERROR] RTSP URL is still empty, retrying in 5 seconds")
time.sleep(5) time.sleep(5)
continue continue
@@ -89,12 +92,12 @@ class RTSPCaptureWorker(threading.Thread):
# cap.set(cv2.CAP_PROP_HW_ACCELERATION, cv2.VIDEO_ACCELERATION_ANY) # cap.set(cv2.CAP_PROP_HW_ACCELERATION, cv2.VIDEO_ACCELERATION_ANY)
if not cap.isOpened(): if not cap.isOpened():
print(f"[ERROR] Cannot open RTSP: {self.rtsp_url}") logger.error(f"[ERROR] Cannot open RTSP: {self.rtsp_url}")
time.sleep(2) time.sleep(2)
self.reconnect_count += 1 self.reconnect_count += 1
continue continue
print(f"[INFO] Successfully opened RTSP: {self.name}") logger.info(f"[INFO] Successfully opened RTSP: {self.name}")
self.reconnect_count = 0 # 重置重连计数 self.reconnect_count = 0 # 重置重连计数
# # 设置帧率（可选） # # 设置帧率（可选）
@@ -104,7 +107,7 @@ class RTSPCaptureWorker(threading.Thread):
ret, frame = cap.read() ret, frame = cap.read()
if not ret: if not ret:
# 检查流是否结束 # 检查流是否结束
print(f"[WARN] Failed to read frame from {self.camera_cfg.name}") logger.warning(f"[WARN] Failed to read frame from {self.camera_cfg.name}")
# 检查是否还有数据 # 检查是否还有数据
time.sleep(0.1) time.sleep(0.1)
@@ -130,7 +133,7 @@ class RTSPCaptureWorker(threading.Thread):
self.raw_queue.put(item, timeout=0.5) self.raw_queue.put(item, timeout=0.5)
except queue.Full: except queue.Full:
print(f"[WARN] Queue full, dropping frame from {self.camera_cfg.name}") logger.warning(f"[WARN] Queue full, dropping frame from {self.camera_cfg.name}")
continue continue
# 控制读取速度，避免过快 # 控制读取速度，避免过快
@@ -139,12 +142,12 @@ class RTSPCaptureWorker(threading.Thread):
cap.release() cap.release()
except Exception as e: except Exception as e:
print(f"[ERROR] Error in RTSP capture for {self.camera_cfg.name}: {e}") logger.error(f"[ERROR] Error in RTSP capture for {self.camera_cfg.name}: {e}")
time.sleep(2) time.sleep(2)
self.reconnect_count += 1 self.reconnect_count += 1
if self.reconnect_count >= self.max_reconnects: if self.reconnect_count >= self.max_reconnects:
print(f"[ERROR] Max reconnects reached for {self.camera_cfg.name}, stopping.") logger.error(f"[ERROR] Max reconnects reached for {self.camera_cfg.name}, stopping.")
def refresh_video_url(self): def refresh_video_url(self):
""" """
@@ -164,14 +167,14 @@ class RTSPCaptureWorker(threading.Thread):
# 解析结果（与test_cam.py相同） # 解析结果（与test_cam.py相同）
if 'data' in result and 'url' in result['data']: if 'data' in result and 'url' in result['data']:
new_url = result['data']['url'] new_url = result['data']['url']
print(f"[INFO] get rtsp url success, URL: {new_url}") logger.info(f"[INFO] get rtsp url success, URL: {new_url}")
return new_url return new_url
else: else:
print(f"[ERROR] get rtsp url failed: {result}") logger.error(f"[ERROR] get rtsp url failed: {result}")
return None return None
except Exception as e: except Exception as e:
print(f"[ERROR] get rtsp url error: {str(e)}") logger.error(f"[ERROR] get rtsp url error: {str(e)}")
return None return None
@@ -199,7 +202,7 @@ class RTSPService:
w = RTSPCaptureWorker(cam, self.raw_queue, self.stop_event) w = RTSPCaptureWorker(cam, self.raw_queue, self.stop_event)
w.start() w.start()
self.capture_workers.append(w) self.capture_workers.append(w)
print("[INFO] Kadian RTSP Service started") logger.info("[INFO] Kadian RTSP Service started")
def stop(self): def stop(self):
self.stop_event.set() self.stop_event.set()
@@ -209,7 +212,7 @@ class RTSPService:
w.join(timeout=2.0) w.join(timeout=2.0)
self.processor.join(timeout=2.0) self.processor.join(timeout=2.0)
self.ws_sender.join(timeout=2.0) self.ws_sender.join(timeout=2.0)
print("[INFO] Service stopped") logger.info("[INFO] Service stopped")
if __name__ == "__main__": if __name__ == "__main__":

9
utils/logger.py
View File

@@ -5,7 +5,7 @@
import logging import logging
import sys import sys
from typing import Optional from typing import Optional
from logging.handlers import RotatingFileHandler from logging.handlers import TimedRotatingFileHandler
from config import settings from config import settings
@@ -56,10 +56,11 @@ def setup_logger(
if log_file or settings.LOG_FILE: if log_file or settings.LOG_FILE:
file_path = log_file or settings.LOG_FILE file_path = log_file or settings.LOG_FILE
try: try:
file_handler = RotatingFileHandler( file_handler = TimedRotatingFileHandler(
file_path, file_path,
maxBytes=10 * 1024 * 1024, # 10MB when='midnight', # 每天午夜轮转
backupCount=5, interval=1, # 间隔1天
backupCount=0, # 保留30天的日志
encoding='utf-8' encoding='utf-8'
) )
file_handler.setLevel(log_level) file_handler.setLevel(log_level)

5
utils/web_socket_sender.py
View File

@@ -5,6 +5,9 @@ import websockets
import threading import threading
import queue import queue
from utils.logger import get_logger
logger = get_logger(__name__)
# ========================= WebSocket 服务线程 ========================= # ========================= WebSocket 服务线程 =========================
class WebSocketSender(threading.Thread): class WebSocketSender(threading.Thread):
def __init__(self, send_queue: queue.Queue, stop_event: threading.Event, ws_host: str, ws_port: int): def __init__(self, send_queue: queue.Queue, stop_event: threading.Event, ws_host: str, ws_port: int):
@@ -42,7 +45,7 @@ class WebSocketSender(threading.Thread):
async def _run_async(self): async def _run_async(self):
async with websockets.serve(self._ws_handler, self.ws_host, self.ws_port): async with websockets.serve(self._ws_handler, self.ws_host, self.ws_port):
print(f"[INFO] WebSocket server started at ws://{self.ws_host}:{self.ws_port}") logger.info(f"[INFO] WebSocket server started at ws://{self.ws_host}:{self.ws_port}")
await self._broadcaster() await self._broadcaster()
def run(self): def run(self):