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 time
from utils.logger import get_logger
logger = get_logger(__name__)
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114)):
shape = img.shape[:2] # h, w
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)
actual_providers = self.session.get_providers()
print("YOLO Providers:", actual_providers)
logger.info("YOLO Providers:", actual_providers)
if "CANNExecutionProvider" in actual_providers:
print("[INFO] YOLO 使用 CANNExecutionProvider（昇腾 NPU）")
logger.info("[INFO] YOLO 使用 CANNExecutionProvider（昇腾 NPU）")
elif 'CUDAExecutionProvider' in actual_providers:
print("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
logger.info("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
else:
print("[INFO] YOLO 使用 CPUExecutionProvider")
logger.info("[INFO] YOLO 使用 CPUExecutionProvider")
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
self.input_name = self.session.get_inputs()[0].name
self.input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
print(f"模型输入名称: {self.input_name}")
print(f"模型输入形状: {self.session.get_inputs()[0].shape}")
print(f"模型输出形状: {self.session.get_outputs()[0].shape}")
logger.info(f"模型输入名称: {self.input_name}")
logger.info(f"模型输入形状: {self.session.get_inputs()[0].shape}")
logger.info(f"模型输出形状: {self.session.get_outputs()[0].shape}")
def preprocess(self, img):
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 onnxruntime as ort
from utils.logger import get_logger
logger = get_logger(__name__)
# -------------------------------------------------
@@ -89,14 +90,14 @@ class YOLOv8_Pose_ONNX:
# 获取真实工作 provider
actual_providers = self.session.get_providers()
print("YOLO Providers:", actual_providers)
logger.info("YOLO Providers:", actual_providers)
if "CANNExecutionProvider" in actual_providers:
print("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）")
logger.info("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）")
elif 'CUDAExecutionProvider' in actual_providers:
print("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
logger.info("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
else:
print("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）")
logger.info("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）")
self.conf_threshold = conf_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_size = (input_size, input_size)
print(f"模型输入名称: {self.input_name}")
print(f"模型输入形状: {self.session.get_inputs()[0].shape}")
print(f"模型输出形状: {self.session.get_outputs()[0].shape}")
logger.info(f"模型输入名称: {self.input_name}")
logger.info(f"模型输入形状: {self.session.get_inputs()[0].shape}")
logger.info(f"模型输出形状: {self.session.get_outputs()[0].shape}")
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 utils.web_socket_sender import WebSocketSender
from utils.logger import get_logger
logger = get_logger(__name__)
WS_HOST = "0.0.0.0"
WS_PORT = 8765
@@ -49,22 +52,22 @@ class RTSPCaptureWorker(threading.Thread):
try:
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
new_url = self.refresh_video_url()
if new_url:
self.rtsp_url = new_url
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，如果是则重新获取
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()
if new_url:
self.rtsp_url = new_url
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)
continue
@@ -89,12 +92,12 @@ class RTSPCaptureWorker(threading.Thread):
# cap.set(cv2.CAP_PROP_HW_ACCELERATION, cv2.VIDEO_ACCELERATION_ANY)
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)
self.reconnect_count += 1
continue
print(f"[INFO] Successfully opened RTSP: {self.name}")
logger.info(f"[INFO] Successfully opened RTSP: {self.name}")
self.reconnect_count = 0 # 重置重连计数
# # 设置帧率（可选）
@@ -104,7 +107,7 @@ class RTSPCaptureWorker(threading.Thread):
ret, frame = cap.read()
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)
@@ -130,7 +133,7 @@ class RTSPCaptureWorker(threading.Thread):
self.raw_queue.put(item, timeout=0.5)
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
# 控制读取速度，避免过快
@@ -139,12 +142,12 @@ class RTSPCaptureWorker(threading.Thread):
cap.release()
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)
self.reconnect_count += 1
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):
"""
@@ -164,14 +167,14 @@ class RTSPCaptureWorker(threading.Thread):
# 解析结果（与test_cam.py相同）
if 'data' in result and 'url' in result['data']:
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
else:
print(f"[ERROR] get rtsp url failed: {result}")
logger.error(f"[ERROR] get rtsp url failed: {result}")
return None
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
@@ -199,7 +202,7 @@ class RTSPService:
w = RTSPCaptureWorker(cam, self.raw_queue, self.stop_event)
w.start()
self.capture_workers.append(w)
print("[INFO] Kadian RTSP Service started")
logger.info("[INFO] Kadian RTSP Service started")
def stop(self):
self.stop_event.set()
@@ -209,7 +212,7 @@ class RTSPService:
w.join(timeout=2.0)
self.processor.join(timeout=2.0)
self.ws_sender.join(timeout=2.0)
print("[INFO] Service stopped")
logger.info("[INFO] Service stopped")
if __name__ == "__main__":

9
utils/logger.py
View File

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

5
utils/web_socket_sender.py
View File

@@ -5,6 +5,9 @@ import websockets
import threading
import queue
from utils.logger import get_logger
logger = get_logger(__name__)
# ========================= WebSocket 服务线程 =========================
class WebSocketSender(threading.Thread):
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 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()
def run(self):