374 lines
12 KiB
Python
374 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import shutil
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
import cv2
|
|
|
|
from app.config import settings
|
|
from app.pipeline import BirdCascadePipeline
|
|
from app.utils import (
|
|
draw_results,
|
|
now_timestamp_minute,
|
|
safe_stem,
|
|
save_json,
|
|
write_video_xml,
|
|
)
|
|
|
|
|
|
def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]:
|
|
"""
|
|
创建单次视频识别结果目录。
|
|
|
|
目录名：
|
|
时间戳_原文件名stem
|
|
|
|
例如：
|
|
runs/video/202605181431_DJI_20260423155914_0006_Z/
|
|
|
|
返回：
|
|
result_dir, result_base_name, original_stem
|
|
"""
|
|
original_stem = safe_stem(original_filename)
|
|
base_name = f"{timestamp}_{original_stem}"
|
|
|
|
result_dir = parent_dir / base_name
|
|
|
|
if not result_dir.exists():
|
|
result_dir.mkdir(parents=True, exist_ok=True)
|
|
return result_dir, base_name, original_stem
|
|
|
|
idx = 1
|
|
while True:
|
|
candidate_base = f"{base_name}_{idx:02d}"
|
|
candidate_dir = parent_dir / candidate_base
|
|
if not candidate_dir.exists():
|
|
candidate_dir.mkdir(parents=True, exist_ok=True)
|
|
return candidate_dir, candidate_base, original_stem
|
|
idx += 1
|
|
|
|
|
|
def iou_xyxy(a: List[float], b: List[float]) -> float:
|
|
ax1, ay1, ax2, ay2 = a
|
|
bx1, by1, bx2, by2 = b
|
|
|
|
x1 = max(ax1, bx1)
|
|
y1 = max(ay1, by1)
|
|
x2 = min(ax2, bx2)
|
|
y2 = min(ay2, by2)
|
|
|
|
inter = max(0.0, x2 - x1) * max(0.0, y2 - y1)
|
|
area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
|
|
area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
|
|
|
|
return inter / (area_a + area_b - inter + 1e-7)
|
|
|
|
|
|
class VideoResultSmoother:
|
|
"""
|
|
降低视频框闪烁：
|
|
1. 非推理帧复用上一推理结果；
|
|
2. 短暂漏检时保留上一结果若干帧；
|
|
3. 对匹配目标的坐标做指数平滑；
|
|
4. 对同一目标累计物种置信度，避免标签和颜色逐帧跳变。
|
|
"""
|
|
|
|
def __init__(self, hold_frames: int = 8, alpha: float = 0.65):
|
|
self.hold_frames = hold_frames
|
|
self.alpha = alpha
|
|
self.last_results: List[Dict] = []
|
|
self.last_update_frame: int = -10**9
|
|
|
|
@staticmethod
|
|
def _public_results(results: List[Dict]) -> List[Dict]:
|
|
return [
|
|
{
|
|
key: value
|
|
for key, value in item.items()
|
|
if not key.startswith("_smoother_")
|
|
}
|
|
for item in results
|
|
]
|
|
|
|
def update(self, current_results: List[Dict], frame_index: int) -> List[Dict]:
|
|
if len(current_results) == 0:
|
|
if frame_index - self.last_update_frame <= self.hold_frames:
|
|
return self._public_results(self.last_results)
|
|
self.last_results = []
|
|
return []
|
|
|
|
if len(self.last_results) == 0:
|
|
self.last_results = []
|
|
for current in current_results:
|
|
item = dict(current)
|
|
class_key = (
|
|
int(item.get("cls_class_id", -1)),
|
|
str(item.get("cls_class_name", "") or ""),
|
|
)
|
|
confidence = max(float(item.get("cls_confidence", 0.0)), 0.01)
|
|
item["_smoother_class_votes"] = {class_key: confidence}
|
|
self.last_results.append(item)
|
|
self.last_update_frame = frame_index
|
|
return self._public_results(self.last_results)
|
|
|
|
smoothed_results = []
|
|
used_last = set()
|
|
|
|
for cur in current_results:
|
|
cur_box = cur["bbox_xyxy"]
|
|
|
|
best_idx = -1
|
|
best_iou = 0.0
|
|
|
|
for idx, prev in enumerate(self.last_results):
|
|
if idx in used_last:
|
|
continue
|
|
|
|
score = iou_xyxy(cur_box, prev["bbox_xyxy"])
|
|
if score > best_iou:
|
|
best_iou = score
|
|
best_idx = idx
|
|
|
|
new_item = dict(cur)
|
|
|
|
if best_idx >= 0 and best_iou >= 0.1:
|
|
prev = self.last_results[best_idx]
|
|
prev_box = prev["bbox_xyxy"]
|
|
|
|
smooth_box = []
|
|
for p, c in zip(prev_box, cur_box):
|
|
smooth_box.append(round(self.alpha * float(p) + (1 - self.alpha) * float(c), 2))
|
|
|
|
new_item["bbox_xyxy"] = smooth_box
|
|
|
|
votes = {
|
|
key: float(value) * 0.90
|
|
for key, value in prev.get(
|
|
"_smoother_class_votes",
|
|
{},
|
|
).items()
|
|
}
|
|
current_key = (
|
|
int(cur.get("cls_class_id", -1)),
|
|
str(cur.get("cls_class_name", "") or ""),
|
|
)
|
|
current_confidence = max(
|
|
float(cur.get("cls_confidence", 0.0)),
|
|
0.01,
|
|
)
|
|
votes[current_key] = votes.get(current_key, 0.0) + current_confidence
|
|
chosen_key = max(votes, key=votes.get)
|
|
new_item["_smoother_class_votes"] = votes
|
|
new_item["cls_class_id"] = chosen_key[0]
|
|
new_item["cls_class_name"] = chosen_key[1]
|
|
if chosen_key != current_key:
|
|
new_item["show_cls_name"] = bool(
|
|
prev.get("show_cls_name", False)
|
|
)
|
|
used_last.add(best_idx)
|
|
else:
|
|
class_key = (
|
|
int(cur.get("cls_class_id", -1)),
|
|
str(cur.get("cls_class_name", "") or ""),
|
|
)
|
|
confidence = max(float(cur.get("cls_confidence", 0.0)), 0.01)
|
|
new_item["_smoother_class_votes"] = {class_key: confidence}
|
|
|
|
smoothed_results.append(new_item)
|
|
|
|
self.last_results = smoothed_results
|
|
self.last_update_frame = frame_index
|
|
|
|
return self._public_results(smoothed_results)
|
|
|
|
def get_for_skipped_frame(self, frame_index: int) -> List[Dict]:
|
|
if frame_index - self.last_update_frame <= self.hold_frames:
|
|
return self._public_results(self.last_results)
|
|
return []
|
|
|
|
|
|
def process_video_file(
|
|
video_path: str,
|
|
pipeline: BirdCascadePipeline,
|
|
original_filename: str,
|
|
sample_interval: Optional[int] = None,
|
|
save_annotated_video: Optional[bool] = None,
|
|
) -> Dict:
|
|
"""
|
|
视频结果自动保存到：
|
|
runs/video/时间戳_原视频名/
|
|
|
|
目录内部生成：
|
|
原视频名.mp4
|
|
时间戳_原视频名.json
|
|
时间戳_原视频名.xml
|
|
原视频名_source.原后缀
|
|
frames/
|
|
|
|
示例：
|
|
runs/video/202605181431_DJI_20260423155914_0006_Z/
|
|
├── DJI_20260423155914_0006_Z.mp4
|
|
├── 202605181431_DJI_20260423155914_0006_Z.json
|
|
├── 202605181431_DJI_20260423155914_0006_Z.xml
|
|
├── DJI_20260423155914_0006_Z_source.MP4
|
|
└── frames/
|
|
"""
|
|
sample_interval = settings.video_sample_interval if sample_interval is None else int(sample_interval)
|
|
save_annotated_video = settings.save_annotated_video if save_annotated_video is None else bool(save_annotated_video)
|
|
|
|
video_root = Path(settings.output_dir) / settings.video_output_subdir
|
|
video_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
timestamp = now_timestamp_minute()
|
|
result_dir, result_base_name, original_stem = make_unique_result_dir(
|
|
parent_dir=video_root,
|
|
timestamp=timestamp,
|
|
original_filename=original_filename,
|
|
)
|
|
|
|
original_suffix = Path(original_filename).suffix or ".mp4"
|
|
|
|
source_video_path = result_dir / f"{original_stem}_source{original_suffix}"
|
|
|
|
# 输出视频固定为 mp4，文件名不带时间戳
|
|
annotated_video_path = result_dir / f"{original_stem}.mp4"
|
|
|
|
# JSON 和 XML 带时间戳
|
|
json_path = result_dir / f"{result_base_name}.json"
|
|
xml_path = result_dir / f"{result_base_name}.xml"
|
|
|
|
frames_dir = result_dir / "frames"
|
|
frames_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
shutil.copy2(video_path, source_video_path)
|
|
|
|
cap = cv2.VideoCapture(str(source_video_path))
|
|
if not cap.isOpened():
|
|
raise ValueError(f"视频打开失败: {source_video_path}")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
if fps is None or fps <= 0:
|
|
fps = 25.0
|
|
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
|
|
|
|
writer = None
|
|
if save_annotated_video:
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
|
writer = cv2.VideoWriter(str(annotated_video_path), fourcc, fps, (width, height))
|
|
|
|
smoother = VideoResultSmoother(
|
|
hold_frames=settings.smooth_hold_frames,
|
|
alpha=settings.box_smooth_alpha,
|
|
)
|
|
|
|
frame_results = []
|
|
processed_frames = 0
|
|
detected_frames = 0
|
|
frame_index = 0
|
|
|
|
start_time = time.time()
|
|
|
|
try:
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
|
|
should_process = frame_index % sample_interval == 0
|
|
|
|
if should_process:
|
|
output = pipeline.predict(frame, return_image=False)
|
|
raw_results = output["results"]
|
|
processed_frames += 1
|
|
|
|
if settings.smooth_video_boxes:
|
|
results = smoother.update(raw_results, frame_index)
|
|
else:
|
|
results = raw_results
|
|
|
|
if len(results) > 0:
|
|
detected_frames += 1
|
|
|
|
frame_results.append(
|
|
{
|
|
"frame_index": frame_index,
|
|
"time_sec": round(frame_index / fps, 3),
|
|
"num_objects": len(results),
|
|
"results": results,
|
|
}
|
|
)
|
|
else:
|
|
if settings.smooth_video_boxes:
|
|
results = smoother.get_for_skipped_frame(frame_index)
|
|
else:
|
|
results = []
|
|
|
|
annotated_frame = draw_results(frame, results)
|
|
|
|
if should_process and len(results) > 0:
|
|
frame_img_path = frames_dir / f"frame_{frame_index:08d}.jpg"
|
|
cv2.imwrite(str(frame_img_path), annotated_frame)
|
|
|
|
if writer is not None:
|
|
writer.write(annotated_frame)
|
|
|
|
frame_index += 1
|
|
|
|
finally:
|
|
cap.release()
|
|
if writer is not None:
|
|
writer.release()
|
|
|
|
elapsed = round(time.time() - start_time, 3)
|
|
|
|
summary = {
|
|
"success": True,
|
|
"type": "video",
|
|
"original_filename": original_filename,
|
|
"result_dir": str(result_dir),
|
|
"source_video": str(source_video_path),
|
|
"annotated_video": str(annotated_video_path) if annotated_video_path.exists() else None,
|
|
"json": str(json_path),
|
|
"xml": str(xml_path),
|
|
"frames_dir": str(frames_dir),
|
|
"video_info": {
|
|
"width": width,
|
|
"height": height,
|
|
"fps": fps,
|
|
"total_frames": total_frames,
|
|
},
|
|
"sample_interval": sample_interval,
|
|
"smooth_video_boxes": settings.smooth_video_boxes,
|
|
"smooth_hold_frames": settings.smooth_hold_frames,
|
|
"box_smooth_alpha": settings.box_smooth_alpha,
|
|
"processed_frames": processed_frames,
|
|
"detected_frames": detected_frames,
|
|
"elapsed_sec": elapsed,
|
|
}
|
|
|
|
json_data = {
|
|
"summary": summary,
|
|
"frames": frame_results,
|
|
}
|
|
|
|
save_json(json_path, json_data)
|
|
|
|
write_video_xml(
|
|
xml_path=xml_path,
|
|
video_filename=annotated_video_path.name,
|
|
video_path=str(annotated_video_path),
|
|
width=width,
|
|
height=height,
|
|
fps=fps,
|
|
total_frames=total_frames,
|
|
sample_interval=sample_interval,
|
|
frame_results=frame_results,
|
|
)
|
|
|
|
return summary
|