|
|
@@ -24,6 +24,7 @@ class PersonTrackingInfo:
|
|
|
bbox: Tuple[int, int, int, int] # (x1, y1, x2, y2)
|
|
|
confidence: float
|
|
|
ptz_position: Optional[Tuple[float, float, int]] = None # (pan, tilt, zoom)
|
|
|
+ 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
|
|
|
|
|
|
@@ -192,16 +193,25 @@ class PairedImageSaver:
|
|
|
# 复制图像避免修改原图
|
|
|
marked_frame = frame.copy()
|
|
|
|
|
|
- # 绘制每个人员的标记
|
|
|
- for i, person in enumerate(persons):
|
|
|
+ # 过滤有效人员(置信度 >= 阈值)
|
|
|
+ person_threshold = 0.8 # 人员检测置信度阈值
|
|
|
+ valid_persons = []
|
|
|
+ for person in persons:
|
|
|
+ conf = person.get('confidence', 0.0)
|
|
|
+ if conf >= person_threshold:
|
|
|
+ valid_persons.append(person)
|
|
|
+
|
|
|
+ # 绘制每个有效人员的标记(使用连续的序号)
|
|
|
+ for i, person in enumerate(valid_persons):
|
|
|
bbox = person.get('bbox', (0, 0, 0, 0))
|
|
|
x1, y1, x2, y2 = bbox
|
|
|
+ conf = person.get('confidence', 0.0)
|
|
|
|
|
|
# 绘制边界框(绿色)
|
|
|
cv2.rectangle(marked_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
|
|
- # 绘制序号标签
|
|
|
- label = f"person_{i}"
|
|
|
+ # 绘制序号标签(带置信度)
|
|
|
+ label = f"person_{i}({conf:.2f})"
|
|
|
(label_w, label_h), baseline = cv2.getTextSize(
|
|
|
label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2
|
|
|
)
|
|
|
@@ -223,12 +233,12 @@ class PairedImageSaver:
|
|
|
(0, 0, 0), 2
|
|
|
)
|
|
|
|
|
|
- # 保存图片
|
|
|
- filename = f"00_panorama_n{len(persons)}.jpg"
|
|
|
+ # 保存图片(使用有效人员数量)
|
|
|
+ filename = f"00_panorama_n{len(valid_persons)}.jpg"
|
|
|
filepath = batch_dir / filename
|
|
|
cv2.imwrite(str(filepath), marked_frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
|
|
|
|
|
|
- logger.info(f"[配对保存] 全景图已保存: {filepath}")
|
|
|
+ logger.info(f"[配对保存] 全景图已保存: {filepath},有效人员 {len(valid_persons)}/{len(persons)}")
|
|
|
return str(filepath)
|
|
|
|
|
|
except Exception as e:
|
|
|
@@ -237,6 +247,7 @@ class PairedImageSaver:
|
|
|
|
|
|
def save_ptz_image(self, batch_id: str, person_index: int,
|
|
|
ptz_frame, ptz_position: Tuple[float, float, int],
|
|
|
+ ptz_bbox: Tuple[int, int, int, int] = None,
|
|
|
person_info: Dict = None) -> Optional[str]:
|
|
|
"""
|
|
|
保存球机聚焦图片
|
|
|
@@ -246,6 +257,7 @@ class PairedImageSaver:
|
|
|
person_index: 人员序号(0-based)
|
|
|
ptz_frame: 球机帧
|
|
|
ptz_position: PTZ位置 (pan, tilt, zoom)
|
|
|
+ ptz_bbox: 球机图中检测到的bbox (x1, y1, x2, y2)
|
|
|
person_info: 额外人员信息
|
|
|
|
|
|
Returns:
|
|
|
@@ -284,6 +296,18 @@ class PairedImageSaver:
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7,
|
|
|
(0, 255, 0), 2
|
|
|
)
|
|
|
+
|
|
|
+ # 绘制PTZ检测到的bbox(红色)
|
|
|
+ if ptz_bbox is not None:
|
|
|
+ x1, y1, x2, y2 = ptz_bbox
|
|
|
+ cv2.rectangle(marked_frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
|
|
|
+ bbox_text = f"PTZ_BBox: ({x1},{y1},{x2},{y2})"
|
|
|
+ cv2.putText(
|
|
|
+ marked_frame, bbox_text,
|
|
|
+ (10, 90),
|
|
|
+ cv2.FONT_HERSHEY_SIMPLEX, 0.6,
|
|
|
+ (0, 0, 255), 2
|
|
|
+ )
|
|
|
|
|
|
# 保存图片
|
|
|
filename = f"01_ptz_person{person_index}_p{int(ptz_position[0])}_t{int(ptz_position[1])}_z{int(ptz_position[2])}.jpg"
|
|
|
@@ -295,6 +319,7 @@ class PairedImageSaver:
|
|
|
# 更新批次信息
|
|
|
if person_index < len(self._current_batch.persons):
|
|
|
self._current_batch.persons[person_index].ptz_position = ptz_position
|
|
|
+ self._current_batch.persons[person_index].ptz_bbox = ptz_bbox
|
|
|
self._current_batch.persons[person_index].ptz_image_saved = True
|
|
|
self._current_batch.persons[person_index].ptz_image_path = str(filepath)
|
|
|
|
|
|
@@ -303,7 +328,7 @@ class PairedImageSaver:
|
|
|
with self._stats_lock:
|
|
|
self._stats['total_ptz_images'] += 1
|
|
|
|
|
|
- logger.info(f"[配对保存] 球机图已保存: {filepath}")
|
|
|
+ logger.info(f"[配对保存] 球机图已保存: {filepath}, BBox={ptz_bbox}")
|
|
|
return str(filepath)
|
|
|
|
|
|
except Exception as e:
|
|
|
@@ -331,8 +356,13 @@ class PairedImageSaver:
|
|
|
f.write(f"\n Person {i}:\n")
|
|
|
f.write(f" Track ID: {person.track_id}\n")
|
|
|
f.write(f" Position: ({person.position[0]:.3f}, {person.position[1]:.3f})\n")
|
|
|
+ f.write(f" BBox: ({person.bbox[0]}, {person.bbox[1]}, {person.bbox[2]}, {person.bbox[3]})\n")
|
|
|
f.write(f" Confidence: {person.confidence:.2f}\n")
|
|
|
f.write(f" PTZ Position: {person.ptz_position}\n")
|
|
|
+ if person.ptz_bbox:
|
|
|
+ f.write(f" PTZ BBox: ({person.ptz_bbox[0]}, {person.ptz_bbox[1]}, {person.ptz_bbox[2]}, {person.ptz_bbox[3]})\n")
|
|
|
+ else:
|
|
|
+ f.write(f" PTZ BBox: None\n")
|
|
|
f.write(f" PTZ Image: {person.ptz_image_path}\n")
|
|
|
|
|
|
logger.info(f"[配对保存] 批次完成: {batch.batch_id}, "
|