|
|
@@ -1,17 +1,19 @@
|
|
|
"""
|
|
|
配对图片保存管理器
|
|
|
将全景检测图片和对应的球机聚焦图片保存到同一目录
|
|
|
+支持 OSS 上传和 batch_info.json 格式
|
|
|
"""
|
|
|
|
|
|
import os
|
|
|
import cv2
|
|
|
import time
|
|
|
+import json
|
|
|
import logging
|
|
|
import threading
|
|
|
from pathlib import Path
|
|
|
from datetime import datetime
|
|
|
-from typing import Optional, List, Dict, Tuple
|
|
|
-from dataclasses import dataclass, field
|
|
|
+from typing import Optional, List, Dict, Tuple, Callable
|
|
|
+from dataclasses import dataclass, field, asdict
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@@ -27,6 +29,7 @@ class PersonInfo:
|
|
|
ptz_bbox: Optional[Tuple[int, int, int, int]] = None # 球机图中检测到的bbox (x1, y1, x2, y2)
|
|
|
ptz_image_saved: bool = False
|
|
|
ptz_image_path: Optional[str] = None
|
|
|
+ ptz_oss_url: Optional[str] = None # 球机图 OSS URL
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
@@ -36,10 +39,13 @@ class DetectionBatch:
|
|
|
timestamp: float
|
|
|
panorama_image: Optional[object] = None # numpy array
|
|
|
panorama_path: Optional[str] = None
|
|
|
+ panorama_oss_url: Optional[str] = None # 全景图 OSS URL
|
|
|
persons: List[PersonInfo] = field(default_factory=list)
|
|
|
total_persons: int = 0
|
|
|
ptz_images_count: int = 0
|
|
|
completed: bool = False
|
|
|
+ device_id: str = '' # 设备编号
|
|
|
+ project_id: str = '' # 项目编号
|
|
|
|
|
|
|
|
|
class PairedImageSaver:
|
|
|
@@ -51,11 +57,16 @@ class PairedImageSaver:
|
|
|
2. 保存全景标记图到批次目录
|
|
|
3. 为每个人员保存对应的球机聚焦图到同一目录
|
|
|
4. 支持时间窗口内的批量保存
|
|
|
+ 5. 支持 OSS 上传
|
|
|
+ 6. 生成 batch_info.json
|
|
|
"""
|
|
|
|
|
|
def __init__(self, base_dir: str = '/home/admin/dsh/paired_images',
|
|
|
time_window: float = 5.0, # 时间窗口(秒)
|
|
|
- max_batches: int = 100):
|
|
|
+ max_batches: int = 100,
|
|
|
+ enable_oss: bool = False,
|
|
|
+ oss_uploader = None,
|
|
|
+ device_config: Dict = None):
|
|
|
"""
|
|
|
初始化
|
|
|
|
|
|
@@ -63,27 +74,39 @@ class PairedImageSaver:
|
|
|
base_dir: 基础保存目录
|
|
|
time_window: 批次时间窗口(秒),同一窗口内的检测归为一批
|
|
|
max_batches: 最大保留批次数量
|
|
|
+ enable_oss: 是否启用 OSS 上传
|
|
|
+ oss_uploader: OSS 上传器实例
|
|
|
+ device_config: 设备配置字典
|
|
|
"""
|
|
|
self.base_dir = Path(base_dir)
|
|
|
self.time_window = time_window
|
|
|
self.max_batches = max_batches
|
|
|
+ self.enable_oss = enable_oss
|
|
|
+ self.oss_uploader = oss_uploader
|
|
|
+ self.device_config = device_config or {}
|
|
|
|
|
|
self._current_batch: Optional[DetectionBatch] = None
|
|
|
self._batch_lock = threading.Lock()
|
|
|
self._last_batch_time = 0.0
|
|
|
|
|
|
+ # 上传状态追踪
|
|
|
+ self._upload_status: Dict[str, Dict] = {} # batch_id -> {panorama: bool, ptz: Dict}
|
|
|
+ self._upload_callback: Optional[Callable] = None
|
|
|
+
|
|
|
# 统计信息
|
|
|
self._stats = {
|
|
|
'total_batches': 0,
|
|
|
'total_persons': 0,
|
|
|
- 'total_ptz_images': 0
|
|
|
+ 'total_ptz_images': 0,
|
|
|
+ 'oss_upload_success': 0,
|
|
|
+ 'oss_upload_failed': 0,
|
|
|
}
|
|
|
self._stats_lock = threading.Lock()
|
|
|
|
|
|
# 确保目录存在
|
|
|
self._ensure_base_dir()
|
|
|
|
|
|
- logger.info(f"[配对保存] 初始化完成: 目录={base_dir}, 时间窗口={time_window}s")
|
|
|
+ logger.info(f"[配对保存] 初始化完成: 目录={base_dir}, 时间窗口={time_window}s, OSS={enable_oss}")
|
|
|
|
|
|
def _ensure_base_dir(self):
|
|
|
"""确保基础目录存在"""
|
|
|
@@ -147,6 +170,10 @@ class PairedImageSaver:
|
|
|
)
|
|
|
person_infos.append(info)
|
|
|
|
|
|
+ # 获取设备信息
|
|
|
+ device_id = self.device_config.get('device_id', 'UNKNOWN')
|
|
|
+ project_id = self.device_config.get('project_id', 'UNKNOWN')
|
|
|
+
|
|
|
# 创建批次记录
|
|
|
self._current_batch = DetectionBatch(
|
|
|
batch_id=batch_id,
|
|
|
@@ -154,9 +181,18 @@ class PairedImageSaver:
|
|
|
panorama_image=panorama_frame,
|
|
|
panorama_path=panorama_path,
|
|
|
persons=person_infos,
|
|
|
- total_persons=len(persons)
|
|
|
+ total_persons=len(persons),
|
|
|
+ device_id=device_id,
|
|
|
+ project_id=project_id
|
|
|
)
|
|
|
|
|
|
+ # 初始化上传状态
|
|
|
+ self._upload_status[batch_id] = {
|
|
|
+ 'panorama': False,
|
|
|
+ 'ptz': {},
|
|
|
+ 'completed': False
|
|
|
+ }
|
|
|
+
|
|
|
self._last_batch_time = current_time
|
|
|
|
|
|
with self._stats_lock:
|
|
|
@@ -168,6 +204,10 @@ class PairedImageSaver:
|
|
|
f"人员={len(persons)}, 目录={batch_dir}"
|
|
|
)
|
|
|
|
|
|
+ # 上传全景图到 OSS
|
|
|
+ if self.enable_oss and panorama_path and self.oss_uploader:
|
|
|
+ self._upload_panorama_to_oss(batch_id, panorama_path)
|
|
|
+
|
|
|
return batch_id
|
|
|
|
|
|
def _save_panorama_image(self, batch_dir: Path, batch_id: str,
|
|
|
@@ -319,6 +359,10 @@ class PairedImageSaver:
|
|
|
self._stats['total_ptz_images'] += 1
|
|
|
|
|
|
logger.info(f"[配对保存] 球机图已保存: {filepath}, BBox={ptz_bbox}")
|
|
|
+
|
|
|
+ # 上传球机图到 OSS
|
|
|
+ if self.enable_oss and self.oss_uploader:
|
|
|
+ self._upload_ptz_to_oss(batch_id, person_index, str(filepath))
|
|
|
else:
|
|
|
# 人员索引超出范围,说明批次信息不一致,跳过保存
|
|
|
logger.warning(f"[配对保存] 人员索引 {person_index} 超出批次范围 {len(self._current_batch.persons)},跳过计数")
|
|
|
@@ -329,21 +373,184 @@ class PairedImageSaver:
|
|
|
logger.error(f"[配对保存] 保存球机图失败: {e}")
|
|
|
return None
|
|
|
|
|
|
+ def _upload_panorama_to_oss(self, batch_id: str, panorama_path: str):
|
|
|
+ """上传全景图到 OSS"""
|
|
|
+ def on_upload_complete(result):
|
|
|
+ if result.success:
|
|
|
+ self._current_batch.panorama_oss_url = result.oss_url
|
|
|
+ self._upload_status[batch_id]['panorama'] = True
|
|
|
+ with self._stats_lock:
|
|
|
+ self._stats['oss_upload_success'] += 1
|
|
|
+ logger.info(f"[OSS] 全景图上传成功: {result.oss_url}")
|
|
|
+ else:
|
|
|
+ with self._stats_lock:
|
|
|
+ self._stats['oss_upload_failed'] += 1
|
|
|
+ logger.error(f"[OSS] 全景图上传失败: {result.error}")
|
|
|
+
|
|
|
+ self.oss_uploader.upload_image(
|
|
|
+ local_path=panorama_path,
|
|
|
+ batch_id=batch_id,
|
|
|
+ image_type='panorama',
|
|
|
+ callback=on_upload_complete
|
|
|
+ )
|
|
|
+
|
|
|
+ def _upload_ptz_to_oss(self, batch_id: str, person_index: int, ptz_path: str):
|
|
|
+ """上传球机图到 OSS"""
|
|
|
+ def on_upload_complete(result):
|
|
|
+ if result.success:
|
|
|
+ # 更新人员信息中的 OSS URL
|
|
|
+ if person_index < len(self._current_batch.persons):
|
|
|
+ self._current_batch.persons[person_index].ptz_oss_url = result.oss_url
|
|
|
+ self._upload_status[batch_id]['ptz'][person_index] = result.oss_url
|
|
|
+ with self._stats_lock:
|
|
|
+ self._stats['oss_upload_success'] += 1
|
|
|
+ logger.info(f"[OSS] 球机图上传成功 (person_{person_index}): {result.oss_url}")
|
|
|
+ else:
|
|
|
+ with self._stats_lock:
|
|
|
+ self._stats['oss_upload_failed'] += 1
|
|
|
+ logger.error(f"[OSS] 球机图上传失败 (person_{person_index}): {result.error}")
|
|
|
+
|
|
|
+ self.oss_uploader.upload_image(
|
|
|
+ local_path=ptz_path,
|
|
|
+ batch_id=batch_id,
|
|
|
+ image_type='ptz',
|
|
|
+ person_index=person_index,
|
|
|
+ callback=on_upload_complete
|
|
|
+ )
|
|
|
+
|
|
|
def _finalize_batch(self, batch: DetectionBatch):
|
|
|
"""完成批次处理"""
|
|
|
batch.completed = True
|
|
|
|
|
|
- # 创建批次信息文件
|
|
|
+ # 等待 OSS 上传完成(最多等待5秒)
|
|
|
+ if self.enable_oss and batch.batch_id in self._upload_status:
|
|
|
+ wait_start = time.time()
|
|
|
+ while time.time() - wait_start < 5.0:
|
|
|
+ status = self._upload_status[batch.batch_id]
|
|
|
+ # 检查全景图是否上传完成
|
|
|
+ panorama_done = status.get('panorama', False) or not batch.panorama_path
|
|
|
+ # 检查所有球机图是否上传完成
|
|
|
+ ptz_done = all(
|
|
|
+ idx in status.get('ptz', {})
|
|
|
+ for idx, person in enumerate(batch.persons)
|
|
|
+ if person.ptz_image_saved
|
|
|
+ )
|
|
|
+ if panorama_done and ptz_done:
|
|
|
+ break
|
|
|
+ time.sleep(0.1)
|
|
|
+
|
|
|
+ # 创建 batch_info.json 文件
|
|
|
try:
|
|
|
batch_dir = self.base_dir / f"batch_{batch.batch_id}"
|
|
|
- info_path = batch_dir / "batch_info.txt"
|
|
|
|
|
|
- with open(info_path, 'w', encoding='utf-8') as f:
|
|
|
+ # 构建 JSON 数据
|
|
|
+ batch_info = self._build_batch_info_json(batch)
|
|
|
+
|
|
|
+ # 保存为 JSON 文件
|
|
|
+ json_path = batch_dir / "batch_info.json"
|
|
|
+ with open(json_path, 'w', encoding='utf-8') as f:
|
|
|
+ json.dump(batch_info, f, ensure_ascii=False, indent=2)
|
|
|
+
|
|
|
+ logger.info(f"[配对保存] batch_info.json 已保存: {json_path}")
|
|
|
+
|
|
|
+ # 同时保留 txt 格式用于兼容(可选)
|
|
|
+ txt_path = batch_dir / "batch_info.txt"
|
|
|
+ self._save_batch_info_txt(batch, txt_path)
|
|
|
+
|
|
|
+ # 标记上传完成
|
|
|
+ if batch.batch_id in self._upload_status:
|
|
|
+ self._upload_status[batch.batch_id]['completed'] = True
|
|
|
+
|
|
|
+ # 触发回调
|
|
|
+ if self._upload_callback:
|
|
|
+ try:
|
|
|
+ self._upload_callback(batch_info)
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"[配对保存] 回调执行错误: {e}")
|
|
|
+
|
|
|
+ logger.info(f"[配对保存] 批次完成: {batch.batch_id}, "
|
|
|
+ f"人员={batch.total_persons}, 球机图={batch.ptz_images_count}")
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"[配对保存] 保存批次信息失败: {e}")
|
|
|
+
|
|
|
+ # 清理旧批次
|
|
|
+ self._cleanup_old_batches()
|
|
|
+
|
|
|
+ def _build_batch_info_json(self, batch: DetectionBatch) -> Dict:
|
|
|
+ """
|
|
|
+ 构建 batch_info.json 数据结构
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ Dict: 批次信息字典
|
|
|
+ """
|
|
|
+ # 人员信息列表
|
|
|
+ persons_list = []
|
|
|
+ for person in batch.persons:
|
|
|
+ person_data = {
|
|
|
+ 'person_index': person.person_index,
|
|
|
+ 'position': {
|
|
|
+ 'x': round(person.position[0], 4),
|
|
|
+ 'y': round(person.position[1], 4)
|
|
|
+ },
|
|
|
+ 'bbox': {
|
|
|
+ 'x1': person.bbox[0],
|
|
|
+ 'y1': person.bbox[1],
|
|
|
+ 'x2': person.bbox[2],
|
|
|
+ 'y2': person.bbox[3]
|
|
|
+ },
|
|
|
+ 'confidence': round(person.confidence, 4),
|
|
|
+ 'ptz_position': {
|
|
|
+ 'pan': round(person.ptz_position[0], 2) if person.ptz_position else None,
|
|
|
+ 'tilt': round(person.ptz_position[1], 2) if person.ptz_position else None,
|
|
|
+ 'zoom': person.ptz_position[2] if person.ptz_position else None
|
|
|
+ } if person.ptz_position else None,
|
|
|
+ 'ptz_bbox': {
|
|
|
+ 'x1': person.ptz_bbox[0],
|
|
|
+ 'y1': person.ptz_bbox[1],
|
|
|
+ 'x2': person.ptz_bbox[2],
|
|
|
+ 'y2': person.ptz_bbox[3]
|
|
|
+ } if person.ptz_bbox else None,
|
|
|
+ 'ptz_image_saved': person.ptz_image_saved,
|
|
|
+ 'ptz_image_path': person.ptz_image_path,
|
|
|
+ 'ptz_oss_url': person.ptz_oss_url
|
|
|
+ }
|
|
|
+ persons_list.append(person_data)
|
|
|
+
|
|
|
+ # 构建完整批次信息
|
|
|
+ batch_info = {
|
|
|
+ 'batch_id': batch.batch_id,
|
|
|
+ 'device_id': batch.device_id,
|
|
|
+ 'project_id': batch.project_id,
|
|
|
+ 'timestamp': batch.timestamp,
|
|
|
+ 'datetime': datetime.fromtimestamp(batch.timestamp).isoformat(),
|
|
|
+ 'total_persons': batch.total_persons,
|
|
|
+ 'ptz_images_count': batch.ptz_images_count,
|
|
|
+ 'panorama': {
|
|
|
+ 'local_path': batch.panorama_path,
|
|
|
+ 'oss_url': batch.panorama_oss_url
|
|
|
+ },
|
|
|
+ 'persons': persons_list,
|
|
|
+ 'upload_status': {
|
|
|
+ 'panorama_uploaded': batch.panorama_oss_url is not None,
|
|
|
+ 'all_ptz_uploaded': all(p.ptz_oss_url is not None for p in batch.persons if p.ptz_image_saved)
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ return batch_info
|
|
|
+
|
|
|
+ def _save_batch_info_txt(self, batch: DetectionBatch, txt_path: Path):
|
|
|
+ """保存批次信息为 TXT 格式(兼容旧版本)"""
|
|
|
+ try:
|
|
|
+ with open(txt_path, 'w', encoding='utf-8') as f:
|
|
|
f.write(f"批次ID: {batch.batch_id}\n")
|
|
|
+ f.write(f"设备ID: {batch.device_id}\n")
|
|
|
+ f.write(f"项目ID: {batch.project_id}\n")
|
|
|
f.write(f"时间戳: {datetime.fromtimestamp(batch.timestamp)}\n")
|
|
|
f.write(f"总人数: {batch.total_persons}\n")
|
|
|
f.write(f"球机图数量: {batch.ptz_images_count}\n")
|
|
|
f.write(f"全景图: {batch.panorama_path}\n")
|
|
|
+ f.write(f"全景图OSS: {batch.panorama_oss_url}\n")
|
|
|
f.write("\n人员详情:\n")
|
|
|
|
|
|
for i, person in enumerate(batch.persons):
|
|
|
@@ -358,15 +565,9 @@ class PairedImageSaver:
|
|
|
else:
|
|
|
f.write(f" PTZ BBox: None\n")
|
|
|
f.write(f" PTZ Image: {person.ptz_image_path}\n")
|
|
|
-
|
|
|
- logger.info(f"[配对保存] 批次完成: {batch.batch_id}, "
|
|
|
- f"人员={batch.total_persons}, 球机图={batch.ptz_images_count}")
|
|
|
-
|
|
|
+ f.write(f" PTZ OSS URL: {person.ptz_oss_url}\n")
|
|
|
except Exception as e:
|
|
|
- logger.error(f"[配对保存] 保存批次信息失败: {e}")
|
|
|
-
|
|
|
- # 清理旧批次
|
|
|
- self._cleanup_old_batches()
|
|
|
+ logger.error(f"[配对保存] 保存 TXT 批次信息失败: {e}")
|
|
|
|
|
|
def _cleanup_old_batches(self):
|
|
|
"""清理旧批次目录"""
|
|
|
@@ -396,6 +597,37 @@ class PairedImageSaver:
|
|
|
with self._stats_lock:
|
|
|
return self._stats.copy()
|
|
|
|
|
|
+ def set_upload_callback(self, callback: Callable):
|
|
|
+ """
|
|
|
+ 设置批次完成回调函数
|
|
|
+
|
|
|
+ Args:
|
|
|
+ callback: 回调函数,接收 batch_info_dict 参数
|
|
|
+ """
|
|
|
+ self._upload_callback = callback
|
|
|
+
|
|
|
+ def get_batch_info(self, batch_id: str) -> Optional[Dict]:
|
|
|
+ """
|
|
|
+ 获取指定批次的 batch_info.json 数据
|
|
|
+
|
|
|
+ Args:
|
|
|
+ batch_id: 批次ID
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ Dict 或 None
|
|
|
+ """
|
|
|
+ try:
|
|
|
+ batch_dir = self.base_dir / f"batch_{batch_id}"
|
|
|
+ json_path = batch_dir / "batch_info.json"
|
|
|
+
|
|
|
+ if json_path.exists():
|
|
|
+ with open(json_path, 'r', encoding='utf-8') as f:
|
|
|
+ return json.load(f)
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f"[配对保存] 读取 batch_info.json 失败: {e}")
|
|
|
+
|
|
|
+ return None
|
|
|
+
|
|
|
def close(self):
|
|
|
"""关闭管理器,完成当前批次"""
|
|
|
with self._batch_lock:
|
|
|
@@ -410,13 +642,18 @@ class PairedImageSaver:
|
|
|
_paired_saver_instance: Optional[PairedImageSaver] = None
|
|
|
|
|
|
|
|
|
-def get_paired_saver(base_dir: str = None, time_window: float = 5.0) -> PairedImageSaver:
|
|
|
+def get_paired_saver(base_dir: str = None, time_window: float = 5.0,
|
|
|
+ enable_oss: bool = False, oss_uploader = None,
|
|
|
+ device_config: Dict = None) -> PairedImageSaver:
|
|
|
"""
|
|
|
获取配对保存管理器实例(单例模式)
|
|
|
|
|
|
Args:
|
|
|
base_dir: 基础保存目录
|
|
|
time_window: 时间窗口
|
|
|
+ enable_oss: 是否启用 OSS 上传
|
|
|
+ oss_uploader: OSS 上传器实例
|
|
|
+ device_config: 设备配置字典
|
|
|
|
|
|
Returns:
|
|
|
PairedImageSaver 实例
|
|
|
@@ -426,7 +663,10 @@ def get_paired_saver(base_dir: str = None, time_window: float = 5.0) -> PairedIm
|
|
|
if _paired_saver_instance is None:
|
|
|
_paired_saver_instance = PairedImageSaver(
|
|
|
base_dir=base_dir or '/home/admin/dsh/paired_images',
|
|
|
- time_window=time_window
|
|
|
+ time_window=time_window,
|
|
|
+ enable_oss=enable_oss,
|
|
|
+ oss_uploader=oss_uploader,
|
|
|
+ device_config=device_config
|
|
|
)
|
|
|
|
|
|
return _paired_saver_instance
|