16337
f70e6b6003
- 新增 core/screenshot_handler.py：监听云端 Redis Stream 截图请求， 抓帧后直传 COS，将结果 URL 写回 Redis - main.py 集成 ScreenshotHandler 的初始化和停止 - requirements.txt 添加 cos-python-sdk-v5 依赖 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
283 lines
9.8 KiB
Python
283 lines
9.8 KiB
Python
"""
|
||
截图处理模块 - 监听云端 Redis Stream 截图请求，截图后上传 COS
|
||
|
||
数据流：
|
||
WVP → XADD edge_snap_request → Edge 消费 → capture frame → upload COS → SET snap:result:{id}
|
||
"""
|
||
|
||
import io
|
||
import json
|
||
import logging
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from typing import Optional
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
from config.settings import get_settings, COSConfig
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Redis 常量
|
||
SNAP_REQUEST_STREAM = "edge_snap_request"
|
||
SNAP_CONSUMER_GROUP = "edge_snap_group"
|
||
SNAP_CONSUMER_NAME = "edge_snap_worker"
|
||
SNAP_RESULT_KEY_PREFIX = "snap:result:"
|
||
SNAP_CACHE_KEY_PREFIX = "snap:cache:"
|
||
SNAP_RESULT_TTL = 60 # 结果 key 60s 过期
|
||
SNAP_CACHE_TTL = 300 # 缓存 key 5min 过期
|
||
|
||
|
||
class ScreenshotHandler:
|
||
"""截图处理器
|
||
|
||
从云端 Redis Stream 消费截图请求，通过 MultiStreamManager 获取最新帧，
|
||
JPEG 编码后直传 COS，最后将结果 URL 写回 Redis。
|
||
"""
|
||
|
||
def __init__(self, cloud_redis, stream_manager):
|
||
"""
|
||
Args:
|
||
cloud_redis: 云端 Redis 客户端 (redis.Redis)
|
||
stream_manager: MultiStreamManager 实例
|
||
"""
|
||
self._cloud_redis = cloud_redis
|
||
self._stream_manager = stream_manager
|
||
self._settings = get_settings()
|
||
self._cos_client = None
|
||
self._cos_config: COSConfig = self._settings.cos
|
||
|
||
self._thread: Optional[threading.Thread] = None
|
||
self._stop_event = threading.Event()
|
||
|
||
# ==================== 生命周期 ====================
|
||
|
||
def start(self):
|
||
"""启动截图监听线程"""
|
||
if not self._cloud_redis:
|
||
logger.warning("[截图] 云端 Redis 不可用，截图处理器未启动")
|
||
return
|
||
if not self._cos_config.secret_id or not self._cos_config.bucket:
|
||
logger.warning("[截图] COS 配置不完整，截图处理器未启动")
|
||
return
|
||
|
||
self._init_cos_client()
|
||
self._ensure_consumer_group()
|
||
|
||
self._stop_event.clear()
|
||
self._thread = threading.Thread(
|
||
target=self._listen_loop,
|
||
name="ScreenshotHandler",
|
||
daemon=True,
|
||
)
|
||
self._thread.start()
|
||
logger.info("[截图] 截图处理器已启动")
|
||
|
||
def stop(self):
|
||
"""停止截图监听线程"""
|
||
self._stop_event.set()
|
||
if self._thread and self._thread.is_alive():
|
||
self._thread.join(timeout=5)
|
||
logger.info("[截图] 截图处理器已停止")
|
||
|
||
# ==================== COS 初始化 ====================
|
||
|
||
def _init_cos_client(self):
|
||
"""初始化 COS 客户端"""
|
||
try:
|
||
from qcloud_cos import CosConfig, CosS3Client
|
||
|
||
config = CosConfig(
|
||
Region=self._cos_config.region,
|
||
SecretId=self._cos_config.secret_id,
|
||
SecretKey=self._cos_config.secret_key,
|
||
)
|
||
self._cos_client = CosS3Client(config)
|
||
logger.info("[截图] COS 客户端初始化成功: bucket=%s, region=%s",
|
||
self._cos_config.bucket, self._cos_config.region)
|
||
except Exception as e:
|
||
logger.error("[截图] COS 客户端初始化失败: %s", e)
|
||
|
||
# ==================== Consumer Group ====================
|
||
|
||
def _ensure_consumer_group(self):
|
||
"""确保 Redis Consumer Group 存在"""
|
||
try:
|
||
self._cloud_redis.xgroup_create(
|
||
SNAP_REQUEST_STREAM, SNAP_CONSUMER_GROUP, id="0", mkstream=True
|
||
)
|
||
logger.info("[截图] 创建 consumer group: %s", SNAP_CONSUMER_GROUP)
|
||
except Exception as e:
|
||
# BUSYGROUP = group already exists, safe to ignore
|
||
if "BUSYGROUP" in str(e):
|
||
logger.debug("[截图] consumer group 已存在")
|
||
else:
|
||
logger.error("[截图] 创建 consumer group 失败: %s", e)
|
||
|
||
# ==================== 主循环 ====================
|
||
|
||
def _listen_loop(self):
|
||
"""XREADGROUP 循环消费截图请求"""
|
||
backoff = 5
|
||
max_backoff = 60
|
||
|
||
while not self._stop_event.is_set():
|
||
try:
|
||
results = self._cloud_redis.xreadgroup(
|
||
SNAP_CONSUMER_GROUP,
|
||
SNAP_CONSUMER_NAME,
|
||
{SNAP_REQUEST_STREAM: ">"},
|
||
count=1,
|
||
block=5000,
|
||
)
|
||
if not results:
|
||
continue
|
||
|
||
backoff = 5 # 重置退避
|
||
|
||
for stream_name, messages in results:
|
||
for msg_id, fields in messages:
|
||
try:
|
||
self._handle_request(fields)
|
||
except Exception as e:
|
||
logger.error("[截图] 处理请求失败: %s", e)
|
||
finally:
|
||
# ACK 消息
|
||
try:
|
||
self._cloud_redis.xack(
|
||
SNAP_REQUEST_STREAM, SNAP_CONSUMER_GROUP, msg_id
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
except Exception as e:
|
||
if self._stop_event.is_set():
|
||
return
|
||
logger.warning("[截图] 监听异常: %s, %ds 后重试", e, backoff)
|
||
self._stop_event.wait(backoff)
|
||
backoff = min(backoff * 2, max_backoff)
|
||
|
||
# ==================== 请求处理 ====================
|
||
|
||
def _handle_request(self, fields: dict):
|
||
"""处理单个截图请求"""
|
||
request_id = fields.get("request_id", "")
|
||
camera_code = fields.get("camera_code", "")
|
||
cos_path = fields.get("cos_path", "")
|
||
|
||
if not request_id or not camera_code or not cos_path:
|
||
logger.warning("[截图] 请求字段不完整: %s", fields)
|
||
self._write_result(request_id, {"status": "error", "message": "请求字段不完整"})
|
||
return
|
||
|
||
logger.info("[截图] 收到截图请求: request_id=%s, camera=%s", request_id, camera_code)
|
||
|
||
# 1. 抓帧
|
||
jpeg_bytes = self._capture_frame(camera_code)
|
||
if jpeg_bytes is None:
|
||
self._write_result(request_id, {
|
||
"status": "error",
|
||
"message": f"摄像头流未连接: {camera_code}"
|
||
})
|
||
return
|
||
|
||
# 2. 上传 COS（失败重试 1 次）
|
||
url = self._upload_to_cos(jpeg_bytes, cos_path)
|
||
if url is None:
|
||
# 重试一次
|
||
logger.info("[截图] COS 上传失败，重试一次...")
|
||
url = self._upload_to_cos(jpeg_bytes, cos_path)
|
||
|
||
if url is None:
|
||
self._write_result(request_id, {
|
||
"status": "error",
|
||
"message": "COS 上传失败"
|
||
})
|
||
return
|
||
|
||
# 3. 写结果
|
||
result = {"status": "ok", "url": url}
|
||
self._write_result(request_id, result)
|
||
|
||
# 4. 写缓存
|
||
cache_key = SNAP_CACHE_KEY_PREFIX + camera_code
|
||
cache_data = json.dumps({"url": url, "timestamp": int(time.time())})
|
||
try:
|
||
self._cloud_redis.set(cache_key, cache_data, ex=SNAP_CACHE_TTL)
|
||
except Exception as e:
|
||
logger.warning("[截图] 写入缓存失败: %s", e)
|
||
|
||
logger.info("[截图] 截图完成: camera=%s, url=%s", camera_code, url)
|
||
|
||
# ==================== 抓帧 ====================
|
||
|
||
def _capture_frame(self, camera_code: str) -> Optional[bytes]:
|
||
"""从 MultiStreamManager 获取最新帧并编码为 JPEG"""
|
||
stream = self._stream_manager.get_stream(camera_code)
|
||
if stream is None:
|
||
logger.warning("[截图] 未找到摄像头流: %s", camera_code)
|
||
return None
|
||
|
||
if not stream.is_connected:
|
||
logger.warning("[截图] 摄像头未连接: %s", camera_code)
|
||
return None
|
||
|
||
frame = stream.get_latest_frame(timeout=2.0)
|
||
if frame is None:
|
||
# 回退：直接从缓冲区读一帧
|
||
frame = stream.read(timeout=2.0)
|
||
|
||
if frame is None:
|
||
logger.warning("[截图] 获取帧超时: %s", camera_code)
|
||
return None
|
||
|
||
# JPEG 编码
|
||
try:
|
||
encode_params = [cv2.IMWRITE_JPEG_QUALITY, 85]
|
||
success, buffer = cv2.imencode(".jpg", frame.image, encode_params)
|
||
if not success:
|
||
logger.error("[截图] JPEG 编码失败: %s", camera_code)
|
||
return None
|
||
return buffer.tobytes()
|
||
except Exception as e:
|
||
logger.error("[截图] 帧编码异常: %s", e)
|
||
return None
|
||
|
||
# ==================== COS 上传 ====================
|
||
|
||
def _upload_to_cos(self, jpeg_bytes: bytes, cos_path: str) -> Optional[str]:
|
||
"""上传 JPEG 到 COS，返回访问 URL"""
|
||
if not self._cos_client:
|
||
logger.error("[截图] COS 客户端未初始化")
|
||
return None
|
||
|
||
try:
|
||
body = io.BytesIO(jpeg_bytes)
|
||
self._cos_client.put_object(
|
||
Bucket=self._cos_config.bucket,
|
||
Body=body,
|
||
Key=cos_path,
|
||
ContentType="image/jpeg",
|
||
)
|
||
url = "https://{}.cos.{}.myqcloud.com/{}".format(
|
||
self._cos_config.bucket, self._cos_config.region, cos_path
|
||
)
|
||
return url
|
||
except Exception as e:
|
||
logger.error("[截图] COS 上传失败: %s", e)
|
||
return None
|
||
|
||
# ==================== Redis 结果 ====================
|
||
|
||
def _write_result(self, request_id: str, data: dict):
|
||
"""将结果写入 Redis key（TTL 60s）"""
|
||
if not request_id:
|
||
return
|
||
key = SNAP_RESULT_KEY_PREFIX + request_id
|
||
try:
|
||
self._cloud_redis.set(key, json.dumps(data), ex=SNAP_RESULT_TTL)
|
||
except Exception as e:
|
||
logger.error("[截图] 写入结果失败: %s", e)
|