249 lines
12 KiB
Python
249 lines
12 KiB
Python
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 实例。
|