参数分离完全版
This commit is contained in:
248
data_pusher.py
Normal file
248
data_pusher.py
Normal file
@ -0,0 +1,248 @@
|
||||
import base64
|
||||
import cv2
|
||||
import numpy as np
|
||||
import requests
|
||||
import time
|
||||
import datetime
|
||||
import logging
|
||||
from fastapi import APIRouter, HTTPException, Body # 切换到 APIRouter
|
||||
from pydantic import BaseModel, HttpUrl
|
||||
# import uvicorn # 不再由此文件运行uvicorn
|
||||
from apscheduler.schedulers.background import BackgroundScheduler
|
||||
from typing import Optional # 用于类型提示
|
||||
# 确保 RFDETRDetector 可以被导入，假设 rfdetr_core.py 在同一目录或 PYTHONPATH 中
|
||||
# from rfdetr_core import RFDETRDetector # 在实际使用中取消注释并确保路径正确
|
||||
|
||||
# 配置日志记录
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# app = FastAPI(title="Data Pusher Service", version="1.0.0") # 移除独立的 FastAPI app
|
||||
pusher_router = APIRouter() # 创建一个 APIRouter
|
||||
|
||||
class DataPusher:
|
||||
def __init__(self, detector): # detector: RFDETRDetector
|
||||
if detector is None:
|
||||
logger.error("DataPusher initialized with a None detector. Push functionality will be impaired.")
|
||||
# 仍然创建实例，但功能会受限，_get_data_payload 会处理 detector is None
|
||||
self.detector = detector
|
||||
self.scheduler = BackgroundScheduler(daemon=True)
|
||||
self.push_job_id = "rt_push_job"
|
||||
self.target_url = None
|
||||
if not self.scheduler.running:
|
||||
try:
|
||||
self.scheduler.start()
|
||||
except Exception as e:
|
||||
logger.error(f"Error starting APScheduler in DataPusher: {e}")
|
||||
|
||||
|
||||
def update_detector_instance(self, detector):
|
||||
"""允许在运行时更新检测器实例，例如当主应用切换模型时"""
|
||||
logger.info(f"DataPusher's detector instance is being updated.")
|
||||
self.detector = detector
|
||||
if detector is None:
|
||||
logger.warning("DataPusher's detector instance updated to None.")
|
||||
|
||||
def _get_data_payload(self):
|
||||
"""获取当前的类别计数和最新标注的帧"""
|
||||
if self.detector is None:
|
||||
logger.warning("DataPusher: Detector not available. Cannot get data payload.")
|
||||
return { # 即使检测器不可用，也返回一个结构，包含空数据
|
||||
# "timestamp": time.time(),
|
||||
"category_counts": {},
|
||||
"frame_base64": None,
|
||||
"error": "Detector not available"
|
||||
}
|
||||
|
||||
category_counts = getattr(self.detector, 'category_counts', {})
|
||||
# 如果 detector 存在但没有 last_annotated_frame (例如模型刚加载还没处理第一帧)
|
||||
last_frame_np = getattr(self.detector, 'last_annotated_frame', None)
|
||||
|
||||
frame_base64 = None
|
||||
if last_frame_np is not None and isinstance(last_frame_np, np.ndarray):
|
||||
try:
|
||||
_, buffer = cv2.imencode('.jpg', last_frame_np)
|
||||
frame_base64 = base64.b64encode(buffer).decode('utf-8')
|
||||
except Exception as e:
|
||||
logger.error(f"Error encoding frame to base64: {e}")
|
||||
|
||||
return {
|
||||
# "timestamp": time.time(),
|
||||
"category_counts": category_counts,
|
||||
"frame_base64": frame_base64
|
||||
}
|
||||
|
||||
def _push_data_task(self):
|
||||
"""执行数据推送的任务"""
|
||||
if not self.target_url:
|
||||
# logger.warning("Target URL not set. Skipping push task.") # 减少日志噪音，仅在初次设置时记录
|
||||
return
|
||||
|
||||
payload = self._get_data_payload()
|
||||
# if payload is None: # _get_data_payload 现在总会返回一个字典
|
||||
# logger.warning("No payload to push.")
|
||||
# return
|
||||
|
||||
try:
|
||||
response = requests.post(self.target_url, json=payload, timeout=5)
|
||||
response.raise_for_status()
|
||||
logger.debug(f"Data pushed successfully to {self.target_url}. Status: {response.status_code}") # 改为 debug 级别
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Error pushing data to {self.target_url}: {e}")
|
||||
except Exception as e:
|
||||
logger.error(f"An unexpected error occurred during data push: {e}")
|
||||
|
||||
def setup_push_schedule(self, frequency: float, target_url: str):
|
||||
"""设置或更新推送计划"""
|
||||
if not isinstance(frequency, (int, float)) or frequency <= 0:
|
||||
raise ValueError("Frequency must be a positive number (pushes per second).")
|
||||
|
||||
self.target_url = str(target_url)
|
||||
interval_seconds = 1.0 / frequency
|
||||
|
||||
if not self.scheduler.running: # 确保调度器正在运行
|
||||
try:
|
||||
logger.info("APScheduler was not running. Attempting to start it now.")
|
||||
self.scheduler.start()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start APScheduler in setup_push_schedule: {e}")
|
||||
raise RuntimeError(f"APScheduler could not be started: {e}")
|
||||
|
||||
|
||||
try:
|
||||
if self.scheduler.get_job(self.push_job_id):
|
||||
self.scheduler.remove_job(self.push_job_id)
|
||||
logger.info(f"Removed existing push job: {self.push_job_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error removing existing job: {e}")
|
||||
|
||||
first_run_time = datetime.datetime.now() + datetime.timedelta(seconds=10)
|
||||
self.scheduler.add_job(
|
||||
self._push_data_task,
|
||||
trigger='interval',
|
||||
seconds=interval_seconds,
|
||||
id=self.push_job_id,
|
||||
next_run_time=first_run_time,
|
||||
replace_existing=True
|
||||
)
|
||||
logger.info(f"Push task scheduled to {self.target_url} every {interval_seconds:.2f}s, starting in 10s.")
|
||||
|
||||
def stop_push_schedule(self):
|
||||
"""停止数据推送任务"""
|
||||
if self.scheduler.get_job(self.push_job_id):
|
||||
try:
|
||||
self.scheduler.remove_job(self.push_job_id)
|
||||
logger.info(f"Push job {self.push_job_id} stopped successfully.")
|
||||
self.target_url = None # 清除目标 URL
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping push job {self.push_job_id}: {e}")
|
||||
else:
|
||||
logger.info("No active push job to stop.")
|
||||
|
||||
def shutdown_scheduler(self):
|
||||
"""安全关闭调度器"""
|
||||
if self.scheduler.running:
|
||||
try:
|
||||
self.scheduler.shutdown()
|
||||
logger.info("DataPusher's APScheduler shut down successfully.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error shutting down DataPusher's APScheduler: {e}")
|
||||
|
||||
def push_specific_payload(self, payload: dict):
|
||||
"""推送一个特定的、已格式化的数据负载到配置的 target_url。"""
|
||||
if not self.target_url:
|
||||
logger.warning("DataPusher: Target URL not set. Cannot push specific payload.")
|
||||
return
|
||||
|
||||
if not payload:
|
||||
logger.warning("DataPusher: Received empty payload for specific push. Skipping.")
|
||||
return
|
||||
|
||||
logger.info(f"Attempting to push specific payload to {self.target_url}")
|
||||
try:
|
||||
response = requests.post(self.target_url, json=payload, timeout=10) # Increased timeout for one-off
|
||||
response.raise_for_status()
|
||||
logger.info(f"Specific payload pushed successfully to {self.target_url}. Status: {response.status_code}")
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Error pushing specific payload to {self.target_url}: {e}")
|
||||
except Exception as e:
|
||||
logger.error(f"An unexpected error occurred during specific payload push: {e}")
|
||||
|
||||
# 全局 DataPusher 实例，将由主应用初始化
|
||||
data_pusher_instance: Optional[DataPusher] = None
|
||||
|
||||
# --- FastAPI Request Body Model ---
|
||||
class PushConfigRequest(BaseModel):
|
||||
frequency: float
|
||||
url: HttpUrl
|
||||
|
||||
# --- FastAPI HTTP Endpoint (using APIRouter) ---
|
||||
@pusher_router.post("/setup_push", summary="配置数据推送任务")
|
||||
async def handle_setup_push(config: PushConfigRequest = Body(...)):
|
||||
global data_pusher_instance
|
||||
if data_pusher_instance is None:
|
||||
# 这个错误理论上不应该发生，如果主应用正确初始化了 data_pusher_instance
|
||||
logger.error("CRITICAL: /setup_push called but data_pusher_instance is None. Main app did not initialize it.")
|
||||
raise HTTPException(status_code=503, detail="DataPusher service not available. Initialization may have failed.")
|
||||
|
||||
if config.frequency <= 0: # Pydantic v2 中可以直接用 gt=0
|
||||
raise HTTPException(status_code=400, detail="Invalid frequency value. Must be a positive number.")
|
||||
|
||||
try:
|
||||
data_pusher_instance.setup_push_schedule(config.frequency, str(config.url))
|
||||
return {
|
||||
"message": "Push task configured successfully.",
|
||||
"frequency_hz": config.frequency,
|
||||
"interval_seconds": 1.0 / config.frequency,
|
||||
"target_url": str(config.url),
|
||||
"first_push_delay_seconds": 10
|
||||
}
|
||||
except ValueError as ve:
|
||||
raise HTTPException(status_code=400, detail=str(ve))
|
||||
except RuntimeError as re: # 例如 APScheduler 启动失败
|
||||
logger.error(f"Runtime error during push schedule setup: {re}")
|
||||
raise HTTPException(status_code=500, detail=str(re))
|
||||
except Exception as e:
|
||||
logger.error(f"Error setting up push schedule: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
|
||||
|
||||
@pusher_router.post("/stop_push", summary="停止当前数据推送任务")
|
||||
async def handle_stop_push():
|
||||
global data_pusher_instance
|
||||
if data_pusher_instance is None:
|
||||
logger.error("CRITICAL: /stop_push called but data_pusher_instance is None.")
|
||||
raise HTTPException(status_code=503, detail="DataPusher service not available.")
|
||||
|
||||
try:
|
||||
data_pusher_instance.stop_push_schedule()
|
||||
return {"message": "Push task stopped successfully if it was running."}
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping push schedule: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=f"Internal server error while stopping push: {str(e)}")
|
||||
|
||||
|
||||
# --- Initialization and Shutdown Functions for Main App ---
|
||||
def initialize_data_pusher(detector_instance_param): # Renamed to avoid conflict
|
||||
"""
|
||||
由主应用程序调用以创建和配置 DataPusher 实例。
|
||||
"""
|
||||
global data_pusher_instance
|
||||
if data_pusher_instance is None:
|
||||
logger.info("Initializing DataPusher instance...")
|
||||
data_pusher_instance = DataPusher(detector_instance_param)
|
||||
else:
|
||||
logger.info("DataPusher instance already initialized. Updating detector instance if provided.")
|
||||
data_pusher_instance.update_detector_instance(detector_instance_param)
|
||||
return data_pusher_instance
|
||||
|
||||
def get_data_pusher_instance() -> Optional[DataPusher]:
|
||||
"""获取 DataPusher 实例 (主要用于主应用可能需要访问它的其他方法，如 shutdown)"""
|
||||
return data_pusher_instance
|
||||
|
||||
# 移除 if __name__ == '__main__' 和 run_pusher_service，因为不再独立运行
|
||||
# 示例代码可以移至主应用的文档或测试脚本中。
|
||||
|
||||
# 注意:
|
||||
# RFDETRDetector 实例的生命周期由 api_server.py (current_detector) 管理。
|
||||
# 当 api_server.py 中的模型切换时，需要有一种机制来更新 DataPusher 内部的 detector 引用。
|
||||
# initialize_data_pusher 可以被多次调用 (例如，在模型切换后)，它会更新 DataPusher 持有的 detector 实例。
|
||||
Reference in New Issue
Block a user