paired_image_saver.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727
  1. """
  2. 配对图片保存管理器
  3. 将全景检测图片和对应的球机聚焦图片保存到同一目录
  4. 支持 OSS 上传和 batch_info.json 格式
  5. """
  6. import os
  7. import cv2
  8. import time
  9. import json
  10. import logging
  11. import threading
  12. from pathlib import Path
  13. from datetime import datetime
  14. from typing import Optional, List, Dict, Tuple, Callable
  15. from dataclasses import dataclass, field, asdict
  16. logger = logging.getLogger(__name__)
  17. @dataclass
  18. class PersonInfo:
  19. """人员信息(轨迹追踪已禁用)"""
  20. person_index: int # 人员序号(0-based)
  21. position: Tuple[float, float] # (x_ratio, y_ratio)
  22. bbox: Tuple[int, int, int, int] # (x1, y1, x2, y2)
  23. confidence: float
  24. ptz_position: Optional[Tuple[float, float, int]] = None # (pan, tilt, zoom)
  25. ptz_bbox: Optional[Tuple[int, int, int, int]] = None # 球机图中检测到的bbox (x1, y1, x2, y2)
  26. ptz_image_saved: bool = False
  27. ptz_image_path: Optional[str] = None
  28. ptz_oss_url: Optional[str] = None # 球机图 OSS URL
  29. @dataclass
  30. class DetectionBatch:
  31. """一批检测记录"""
  32. batch_id: str
  33. timestamp: float
  34. panorama_image: Optional[object] = None # numpy array
  35. panorama_path: Optional[str] = None
  36. panorama_oss_url: Optional[str] = None # 全景图 OSS URL
  37. persons: List[PersonInfo] = field(default_factory=list)
  38. total_persons: int = 0
  39. ptz_images_count: int = 0
  40. completed: bool = False
  41. device_id: str = '' # 设备编号
  42. project_id: str = '' # 项目编号
  43. class PairedImageSaver:
  44. """
  45. 配对图片保存管理器
  46. 功能:
  47. 1. 为每次全景检测创建批次目录
  48. 2. 保存全景标记图到批次目录
  49. 3. 为每个人员保存对应的球机聚焦图到同一目录
  50. 4. 支持时间窗口内的批量保存
  51. 5. 支持 OSS 上传
  52. 6. 生成 batch_info.json
  53. """
  54. def __init__(self, base_dir: str = '/home/admin/dsh/paired_images',
  55. time_window: float = 5.0, # 时间窗口(秒)
  56. max_batches: int = 100,
  57. enable_oss: bool = False,
  58. oss_uploader = None,
  59. device_config: Dict = None):
  60. """
  61. 初始化
  62. Args:
  63. base_dir: 基础保存目录
  64. time_window: 批次时间窗口(秒),同一窗口内的检测归为一批
  65. max_batches: 最大保留批次数量
  66. enable_oss: 是否启用 OSS 上传
  67. oss_uploader: OSS 上传器实例
  68. device_config: 设备配置字典
  69. """
  70. self.base_dir = Path(base_dir)
  71. self.time_window = time_window
  72. self.max_batches = max_batches
  73. # 从配置模块读取 OSS 和设备配置(确保即使外部不传也能正确配置)
  74. try:
  75. from config import S3_COMPATIBLE_CONFIG, DEVICE_CONFIG
  76. # OSS 配置:优先使用传入参数,否则从配置模块读取
  77. if enable_oss or (not enable_oss and S3_COMPATIBLE_CONFIG.get('enabled', False)):
  78. self.enable_oss = True
  79. else:
  80. self.enable_oss = enable_oss
  81. # OSS 上传器:优先使用传入的实例,否则从全局获取
  82. if oss_uploader is not None:
  83. self.oss_uploader = oss_uploader
  84. elif self.enable_oss:
  85. try:
  86. from oss_uploader import get_oss_uploader
  87. self.oss_uploader = get_oss_uploader()
  88. if not self.oss_uploader.running:
  89. self.oss_uploader.start()
  90. except Exception as e:
  91. logger.warning(f"[配对保存] OSS 上传器初始化失败: {e}")
  92. self.oss_uploader = None
  93. self.enable_oss = False
  94. else:
  95. self.oss_uploader = None
  96. # 设备配置:合并传入参数和配置模块
  97. self.device_config = DEVICE_CONFIG.copy()
  98. if device_config:
  99. self.device_config.update(device_config)
  100. except ImportError:
  101. # 配置模块不可用时使用传入参数
  102. self.enable_oss = enable_oss
  103. self.oss_uploader = oss_uploader
  104. self.device_config = device_config or {}
  105. self._current_batch: Optional[DetectionBatch] = None
  106. self._batch_lock = threading.Lock()
  107. self._last_batch_time = 0.0
  108. # 上传状态追踪
  109. self._upload_status: Dict[str, Dict] = {} # batch_id -> {panorama: bool, ptz: Dict}
  110. self._upload_callback: Optional[Callable] = None
  111. # 统计信息
  112. self._stats = {
  113. 'total_batches': 0,
  114. 'total_persons': 0,
  115. 'total_ptz_images': 0,
  116. 'oss_upload_success': 0,
  117. 'oss_upload_failed': 0,
  118. }
  119. self._stats_lock = threading.Lock()
  120. # 确保目录存在
  121. self._ensure_base_dir()
  122. logger.info(f"[配对保存] 初始化完成: 目录={base_dir}, 时间窗口={time_window}s, OSS={enable_oss}")
  123. def _ensure_base_dir(self):
  124. """确保基础目录存在"""
  125. try:
  126. self.base_dir.mkdir(parents=True, exist_ok=True)
  127. except Exception as e:
  128. logger.error(f"[配对保存] 创建目录失败: {e}")
  129. def _generate_batch_id(self) -> str:
  130. """生成批次ID"""
  131. return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
  132. def _create_batch_dir(self, batch_id: str) -> Path:
  133. """创建批次目录"""
  134. batch_dir = self.base_dir / f"batch_{batch_id}"
  135. try:
  136. batch_dir.mkdir(parents=True, exist_ok=True)
  137. return batch_dir
  138. except Exception as e:
  139. logger.error(f"[配对保存] 创建批次目录失败: {e}")
  140. return self.base_dir
  141. def start_new_batch(self, panorama_frame, persons: List[Dict]) -> Optional[str]:
  142. """
  143. 开始新批次
  144. Args:
  145. panorama_frame: 全景帧图像
  146. persons: 人员列表,每项包含 track_id, position, bbox, confidence
  147. Returns:
  148. batch_id: 批次ID,失败返回 None
  149. """
  150. with self._batch_lock:
  151. current_time = time.time()
  152. # 完成上一批次(如果有)
  153. # 注意:每次检测都创建独立批次,不复用,确保 batch_info 与实际检测一致
  154. if self._current_batch is not None:
  155. self._finalize_batch(self._current_batch)
  156. # 创建新批次
  157. batch_id = self._generate_batch_id()
  158. batch_dir = self._create_batch_dir(batch_id)
  159. # 保存全景图片
  160. panorama_path = None
  161. if panorama_frame is not None:
  162. panorama_path = self._save_panorama_image(
  163. batch_dir, batch_id, panorama_frame, persons
  164. )
  165. # 创建人员信息(轨迹追踪已禁用,使用序号代替track_id)
  166. person_infos = []
  167. for i, p in enumerate(persons):
  168. info = PersonInfo(
  169. person_index=i,
  170. position=p.get('position', (0, 0)),
  171. bbox=p.get('bbox', (0, 0, 0, 0)),
  172. confidence=p.get('confidence', 0.0)
  173. )
  174. person_infos.append(info)
  175. # 获取设备信息
  176. device_id = self.device_config.get('device_id', 'UNKNOWN')
  177. project_id = self.device_config.get('project_id', 'UNKNOWN')
  178. # 创建批次记录
  179. self._current_batch = DetectionBatch(
  180. batch_id=batch_id,
  181. timestamp=current_time,
  182. panorama_image=panorama_frame,
  183. panorama_path=panorama_path,
  184. persons=person_infos,
  185. total_persons=len(persons),
  186. device_id=device_id,
  187. project_id=project_id
  188. )
  189. # 初始化上传状态
  190. self._upload_status[batch_id] = {
  191. 'panorama': False,
  192. 'ptz': {},
  193. 'completed': False
  194. }
  195. self._last_batch_time = current_time
  196. with self._stats_lock:
  197. self._stats['total_batches'] += 1
  198. self._stats['total_persons'] += len(persons)
  199. logger.info(
  200. f"[配对保存] 新批次创建: {batch_id}, "
  201. f"人员={len(persons)}, 目录={batch_dir}"
  202. )
  203. # 上传全景图到 OSS
  204. if self.enable_oss and panorama_path and self.oss_uploader:
  205. self._upload_panorama_to_oss(batch_id, panorama_path)
  206. return batch_id
  207. def _save_panorama_image(self, batch_dir: Path, batch_id: str,
  208. frame, persons: List[Dict]) -> Optional[str]:
  209. """
  210. 保存全景标记图片
  211. Args:
  212. batch_dir: 批次目录
  213. batch_id: 批次ID
  214. frame: 全景帧
  215. persons: 人员列表(已由调用方过滤,此处不再过滤)
  216. Returns:
  217. 保存路径或 None
  218. """
  219. try:
  220. # 复制图像避免修改原图
  221. marked_frame = frame.copy()
  222. # 绘制每个人员的标记(使用连续的序号)
  223. # 注意:persons 已由调用方(coordinator)过滤,置信度均 >= 阈值
  224. for i, person in enumerate(persons):
  225. bbox = person.get('bbox', (0, 0, 0, 0))
  226. x1, y1, x2, y2 = bbox
  227. conf = person.get('confidence', 0.0)
  228. # 绘制边界框(绿色)
  229. cv2.rectangle(marked_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
  230. # 绘制序号标签(带置信度)
  231. label = f"person_{i}({conf:.2f})"
  232. (label_w, label_h), baseline = cv2.getTextSize(
  233. label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2
  234. )
  235. # 标签背景
  236. cv2.rectangle(
  237. marked_frame,
  238. (x1, y1 - label_h - 8),
  239. (x1 + label_w, y1),
  240. (0, 255, 0),
  241. -1
  242. )
  243. # 标签文字(黑色)
  244. cv2.putText(
  245. marked_frame, label,
  246. (x1, y1 - 4),
  247. cv2.FONT_HERSHEY_SIMPLEX, 0.8,
  248. (0, 0, 0), 2
  249. )
  250. # 保存图片(使用人员数量)
  251. filename = f"00_panorama_n{len(persons)}.png"
  252. filepath = batch_dir / filename
  253. # 保存全景图(PNG无损格式,不压缩)
  254. cv2.imwrite(str(filepath), marked_frame)
  255. logger.info(f"[配对保存] 全景图已保存: {filepath},人员数量 {len(persons)}")
  256. return str(filepath)
  257. except Exception as e:
  258. logger.error(f"[配对保存] 保存全景图失败: {e}")
  259. return None
  260. def save_ptz_image(self, batch_id: str, person_index: int,
  261. ptz_frame, ptz_position: Tuple[float, float, int],
  262. ptz_bbox: Tuple[int, int, int, int] = None,
  263. person_info: Dict = None) -> Optional[str]:
  264. """
  265. 保存球机聚焦图片
  266. Args:
  267. batch_id: 批次ID
  268. person_index: 人员序号(0-based)
  269. ptz_frame: 球机帧
  270. ptz_position: PTZ位置 (pan, tilt, zoom)
  271. ptz_bbox: 球机图中检测到的bbox (x1, y1, x2, y2)
  272. person_info: 额外人员信息
  273. Returns:
  274. 保存路径或 None
  275. """
  276. with self._batch_lock:
  277. if self._current_batch is None or self._current_batch.batch_id != batch_id:
  278. logger.warning(f"[配对保存] 批次不存在或已过期: {batch_id}")
  279. return None
  280. batch_dir = self.base_dir / f"batch_{batch_id}"
  281. try:
  282. # 复制图像
  283. marked_frame = ptz_frame.copy() if ptz_frame is not None else None
  284. if marked_frame is not None:
  285. # 在球机图上添加标记(如果有检测框)
  286. h, w = marked_frame.shape[:2]
  287. # 添加PTZ位置信息到图片
  288. pan, tilt, zoom = ptz_position
  289. info_text = f"PTZ: P={pan:.1f} T={tilt:.1f} Z={zoom}"
  290. cv2.putText(
  291. marked_frame, info_text,
  292. (10, 30),
  293. cv2.FONT_HERSHEY_SIMPLEX, 0.7,
  294. (0, 255, 0), 2
  295. )
  296. # 添加人员序号
  297. person_text = f"person_{person_index}"
  298. cv2.putText(
  299. marked_frame, person_text,
  300. (10, 60),
  301. cv2.FONT_HERSHEY_SIMPLEX, 0.7,
  302. (0, 255, 0), 2
  303. )
  304. # 绘制PTZ检测到的bbox(红色)
  305. if ptz_bbox is not None:
  306. x1, y1, x2, y2 = ptz_bbox
  307. cv2.rectangle(marked_frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
  308. bbox_text = f"PTZ_BBox: ({x1},{y1},{x2},{y2})"
  309. cv2.putText(
  310. marked_frame, bbox_text,
  311. (10, 90),
  312. cv2.FONT_HERSHEY_SIMPLEX, 0.6,
  313. (0, 0, 255), 2
  314. )
  315. # 保存图片
  316. filename = f"01_ptz_person{person_index}_p{int(ptz_position[0])}_t{int(ptz_position[1])}_z{int(ptz_position[2])}.jpg"
  317. filepath = batch_dir / filename
  318. if marked_frame is not None:
  319. cv2.imwrite(str(filepath), marked_frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
  320. # 更新批次信息
  321. # 【关键修复】只有当人员索引有效时才更新和计数
  322. if person_index < len(self._current_batch.persons):
  323. self._current_batch.persons[person_index].ptz_position = ptz_position
  324. self._current_batch.persons[person_index].ptz_bbox = ptz_bbox
  325. self._current_batch.persons[person_index].ptz_image_saved = True
  326. self._current_batch.persons[person_index].ptz_image_path = str(filepath)
  327. self._current_batch.ptz_images_count += 1
  328. with self._stats_lock:
  329. self._stats['total_ptz_images'] += 1
  330. logger.info(f"[配对保存] 球机图已保存: {filepath}, BBox={ptz_bbox}")
  331. # 上传球机图到 OSS
  332. if self.enable_oss and self.oss_uploader:
  333. self._upload_ptz_to_oss(batch_id, person_index, str(filepath))
  334. else:
  335. # 人员索引超出范围,说明批次信息不一致,跳过保存
  336. logger.warning(f"[配对保存] 人员索引 {person_index} 超出批次范围 {len(self._current_batch.persons)},跳过计数")
  337. return str(filepath)
  338. except Exception as e:
  339. logger.error(f"[配对保存] 保存球机图失败: {e}")
  340. return None
  341. def _upload_panorama_to_oss(self, batch_id: str, panorama_path: str):
  342. """上传全景图到 OSS"""
  343. def on_upload_complete(result):
  344. if result.success:
  345. self._current_batch.panorama_oss_url = result.oss_url
  346. self._upload_status[batch_id]['panorama'] = True
  347. with self._stats_lock:
  348. self._stats['oss_upload_success'] += 1
  349. logger.info(f"[OSS] 全景图上传成功: {result.oss_url}")
  350. else:
  351. with self._stats_lock:
  352. self._stats['oss_upload_failed'] += 1
  353. logger.error(f"[OSS] 全景图上传失败: {result.error}")
  354. self.oss_uploader.upload_image(
  355. local_path=panorama_path,
  356. batch_id=batch_id,
  357. image_type='panorama',
  358. callback=on_upload_complete
  359. )
  360. def _upload_ptz_to_oss(self, batch_id: str, person_index: int, ptz_path: str):
  361. """上传球机图到 OSS"""
  362. def on_upload_complete(result):
  363. if result.success:
  364. # 更新人员信息中的 OSS URL
  365. if person_index < len(self._current_batch.persons):
  366. self._current_batch.persons[person_index].ptz_oss_url = result.oss_url
  367. self._upload_status[batch_id]['ptz'][person_index] = result.oss_url
  368. with self._stats_lock:
  369. self._stats['oss_upload_success'] += 1
  370. logger.info(f"[OSS] 球机图上传成功 (person_{person_index}): {result.oss_url}")
  371. else:
  372. with self._stats_lock:
  373. self._stats['oss_upload_failed'] += 1
  374. logger.error(f"[OSS] 球机图上传失败 (person_{person_index}): {result.error}")
  375. self.oss_uploader.upload_image(
  376. local_path=ptz_path,
  377. batch_id=batch_id,
  378. image_type='ptz',
  379. person_index=person_index,
  380. callback=on_upload_complete
  381. )
  382. def _finalize_batch(self, batch: DetectionBatch):
  383. """完成批次处理"""
  384. batch.completed = True
  385. # 等待 OSS 上传完成(最多等待5秒)
  386. if self.enable_oss and batch.batch_id in self._upload_status:
  387. wait_start = time.time()
  388. while time.time() - wait_start < 5.0:
  389. status = self._upload_status[batch.batch_id]
  390. # 检查全景图是否上传完成
  391. panorama_done = status.get('panorama', False) or not batch.panorama_path
  392. # 检查所有球机图是否上传完成
  393. ptz_done = all(
  394. idx in status.get('ptz', {})
  395. for idx, person in enumerate(batch.persons)
  396. if person.ptz_image_saved
  397. )
  398. if panorama_done and ptz_done:
  399. break
  400. time.sleep(0.1)
  401. # 创建 batch_info.json 文件
  402. try:
  403. batch_dir = self.base_dir / f"batch_{batch.batch_id}"
  404. # 构建 JSON 数据
  405. batch_info = self._build_batch_info_json(batch)
  406. # 保存为 JSON 文件
  407. json_path = batch_dir / "batch_info.json"
  408. with open(json_path, 'w', encoding='utf-8') as f:
  409. json.dump(batch_info, f, ensure_ascii=False, indent=2)
  410. logger.info(f"[配对保存] batch_info.json 已保存: {json_path}")
  411. # 同时保留 txt 格式用于兼容(可选)
  412. txt_path = batch_dir / "batch_info.txt"
  413. self._save_batch_info_txt(batch, txt_path)
  414. # 标记上传完成
  415. if batch.batch_id in self._upload_status:
  416. self._upload_status[batch.batch_id]['completed'] = True
  417. # 触发回调
  418. if self._upload_callback:
  419. try:
  420. self._upload_callback(batch_info)
  421. except Exception as e:
  422. logger.error(f"[配对保存] 回调执行错误: {e}")
  423. logger.info(f"[配对保存] 批次完成: {batch.batch_id}, "
  424. f"人员={batch.total_persons}, 球机图={batch.ptz_images_count}")
  425. except Exception as e:
  426. logger.error(f"[配对保存] 保存批次信息失败: {e}")
  427. # 清理旧批次
  428. self._cleanup_old_batches()
  429. def _build_batch_info_json(self, batch: DetectionBatch) -> Dict:
  430. """
  431. 构建 batch_info.json 数据结构
  432. Returns:
  433. Dict: 批次信息字典
  434. """
  435. # 人员信息列表
  436. persons_list = []
  437. for person in batch.persons:
  438. person_data = {
  439. 'person_index': person.person_index,
  440. 'position': {
  441. 'x': round(person.position[0], 4),
  442. 'y': round(person.position[1], 4)
  443. },
  444. 'bbox': {
  445. 'x1': person.bbox[0],
  446. 'y1': person.bbox[1],
  447. 'x2': person.bbox[2],
  448. 'y2': person.bbox[3]
  449. },
  450. 'confidence': round(person.confidence, 4),
  451. 'ptz_position': {
  452. 'pan': round(person.ptz_position[0], 2) if person.ptz_position else None,
  453. 'tilt': round(person.ptz_position[1], 2) if person.ptz_position else None,
  454. 'zoom': person.ptz_position[2] if person.ptz_position else None
  455. } if person.ptz_position else None,
  456. 'ptz_bbox': {
  457. 'x1': person.ptz_bbox[0],
  458. 'y1': person.ptz_bbox[1],
  459. 'x2': person.ptz_bbox[2],
  460. 'y2': person.ptz_bbox[3]
  461. } if person.ptz_bbox else None,
  462. 'ptz_image_saved': person.ptz_image_saved,
  463. 'ptz_image_path': person.ptz_image_path,
  464. 'ptz_oss_url': person.ptz_oss_url
  465. }
  466. persons_list.append(person_data)
  467. # 构建完整批次信息
  468. batch_info = {
  469. 'batch_id': batch.batch_id,
  470. 'device_id': batch.device_id,
  471. 'project_id': batch.project_id,
  472. 'timestamp': batch.timestamp,
  473. 'datetime': datetime.fromtimestamp(batch.timestamp).isoformat(),
  474. 'total_persons': batch.total_persons,
  475. 'ptz_images_count': batch.ptz_images_count,
  476. 'panorama': {
  477. 'local_path': batch.panorama_path,
  478. 'oss_url': batch.panorama_oss_url
  479. },
  480. 'persons': persons_list,
  481. 'upload_status': {
  482. 'panorama_uploaded': batch.panorama_oss_url is not None,
  483. 'all_ptz_uploaded': all(p.ptz_oss_url is not None for p in batch.persons if p.ptz_image_saved)
  484. }
  485. }
  486. return batch_info
  487. def _save_batch_info_txt(self, batch: DetectionBatch, txt_path: Path):
  488. """保存批次信息为 TXT 格式(兼容旧版本)"""
  489. try:
  490. with open(txt_path, 'w', encoding='utf-8') as f:
  491. f.write(f"批次ID: {batch.batch_id}\n")
  492. f.write(f"设备ID: {batch.device_id}\n")
  493. f.write(f"项目ID: {batch.project_id}\n")
  494. f.write(f"时间戳: {datetime.fromtimestamp(batch.timestamp)}\n")
  495. f.write(f"总人数: {batch.total_persons}\n")
  496. f.write(f"球机图数量: {batch.ptz_images_count}\n")
  497. f.write(f"全景图: {batch.panorama_path}\n")
  498. f.write(f"全景图OSS: {batch.panorama_oss_url}\n")
  499. f.write("\n人员详情:\n")
  500. for i, person in enumerate(batch.persons):
  501. f.write(f"\n Person {i}:\n")
  502. f.write(f" Person Index: {person.person_index}\n")
  503. f.write(f" Position: ({person.position[0]:.3f}, {person.position[1]:.3f})\n")
  504. f.write(f" BBox: ({person.bbox[0]}, {person.bbox[1]}, {person.bbox[2]}, {person.bbox[3]})\n")
  505. f.write(f" Confidence: {person.confidence:.2f}\n")
  506. f.write(f" PTZ Position: {person.ptz_position}\n")
  507. if person.ptz_bbox:
  508. f.write(f" PTZ BBox: ({person.ptz_bbox[0]}, {person.ptz_bbox[1]}, {person.ptz_bbox[2]}, {person.ptz_bbox[3]})\n")
  509. else:
  510. f.write(f" PTZ BBox: None\n")
  511. f.write(f" PTZ Image: {person.ptz_image_path}\n")
  512. f.write(f" PTZ OSS URL: {person.ptz_oss_url}\n")
  513. except Exception as e:
  514. logger.error(f"[配对保存] 保存 TXT 批次信息失败: {e}")
  515. def _cleanup_old_batches(self):
  516. """清理旧批次目录"""
  517. try:
  518. batch_dirs = sorted(
  519. [d for d in self.base_dir.iterdir() if d.is_dir() and d.name.startswith('batch_')],
  520. key=lambda x: x.stat().st_mtime
  521. )
  522. if len(batch_dirs) > self.max_batches:
  523. to_delete = batch_dirs[:len(batch_dirs) - self.max_batches]
  524. for d in to_delete:
  525. import shutil
  526. shutil.rmtree(d)
  527. logger.info(f"[配对保存] 清理旧批次: {d.name}")
  528. except Exception as e:
  529. logger.error(f"[配对保存] 清理旧批次失败: {e}")
  530. def get_current_batch_id(self) -> Optional[str]:
  531. """获取当前批次ID"""
  532. with self._batch_lock:
  533. return self._current_batch.batch_id if self._current_batch else None
  534. def get_stats(self) -> Dict:
  535. """获取统计信息"""
  536. with self._stats_lock:
  537. return self._stats.copy()
  538. def set_upload_callback(self, callback: Callable):
  539. """
  540. 设置批次完成回调函数
  541. Args:
  542. callback: 回调函数,接收 batch_info_dict 参数
  543. """
  544. self._upload_callback = callback
  545. def get_batch_info(self, batch_id: str) -> Optional[Dict]:
  546. """
  547. 获取指定批次的 batch_info.json 数据
  548. Args:
  549. batch_id: 批次ID
  550. Returns:
  551. Dict 或 None
  552. """
  553. try:
  554. batch_dir = self.base_dir / f"batch_{batch_id}"
  555. json_path = batch_dir / "batch_info.json"
  556. if json_path.exists():
  557. with open(json_path, 'r', encoding='utf-8') as f:
  558. return json.load(f)
  559. except Exception as e:
  560. logger.error(f"[配对保存] 读取 batch_info.json 失败: {e}")
  561. return None
  562. def close(self):
  563. """关闭管理器,完成当前批次"""
  564. with self._batch_lock:
  565. if self._current_batch is not None:
  566. self._finalize_batch(self._current_batch)
  567. self._current_batch = None
  568. logger.info("[配对保存] 管理器已关闭")
  569. # 全局单例实例
  570. _paired_saver_instance: Optional[PairedImageSaver] = None
  571. def get_paired_saver(base_dir: str = None, time_window: float = 5.0,
  572. enable_oss: bool = False, oss_uploader = None,
  573. device_config: Dict = None) -> PairedImageSaver:
  574. """
  575. 获取配对保存管理器实例(单例模式)
  576. 如果实例已存在但缺少 OSS/设备配置,会自动从配置模块更新
  577. Args:
  578. base_dir: 基础保存目录
  579. time_window: 时间窗口
  580. enable_oss: 是否启用 OSS 上传
  581. oss_uploader: OSS 上传器实例
  582. device_config: 设备配置字典
  583. Returns:
  584. PairedImageSaver 实例
  585. """
  586. global _paired_saver_instance
  587. if _paired_saver_instance is None:
  588. _paired_saver_instance = PairedImageSaver(
  589. base_dir=base_dir or '/home/admin/dsh/paired_images',
  590. time_window=time_window,
  591. enable_oss=enable_oss,
  592. oss_uploader=oss_uploader,
  593. device_config=device_config
  594. )
  595. else:
  596. # 单例已存在,检查是否需要更新配置
  597. # PairedImageSaver.__init__ 已从配置模块自动读取,
  598. # 这里只处理外部显式传入的参数覆盖
  599. if oss_uploader is not None and _paired_saver_instance.oss_uploader is None:
  600. _paired_saver_instance.oss_uploader = oss_uploader
  601. _paired_saver_instance.enable_oss = True
  602. logger.info("[配对保存] 更新 OSS 上传器配置")
  603. if device_config is not None:
  604. _paired_saver_instance.device_config.update(device_config)
  605. logger.info(f"[配对保存] 更新设备配置: {device_config}")
  606. return _paired_saver_instance
  607. def reset_paired_saver():
  608. """重置单例实例(用于测试)"""
  609. global _paired_saver_instance
  610. if _paired_saver_instance is not None:
  611. _paired_saver_instance.close()
  612. _paired_saver_instance = None