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 实例。
|