75 lines
3.3 KiB
Python
75 lines
3.3 KiB
Python
"""Camera 01 录像：直接 YOLO 与共享实时缩图管线的一致性抽样验证。
|
||
|
||
手动运行：python tests/validate_yolo_replay.py <video> <weights> [sample_count] [device]
|
||
"""
|
||
import sys
|
||
|
||
import cv2
|
||
from ultralytics import YOLO
|
||
|
||
|
||
def iou(a, b):
|
||
x1, y1, x2, y2 = max(a[0], b[0]), max(a[1], b[1]), min(a[2], b[2]), min(a[3], b[3])
|
||
inter = max(0, x2 - x1) * max(0, y2 - y1)
|
||
union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
|
||
return inter / union if union > 0 else 0
|
||
|
||
|
||
def boxes(result):
|
||
if result.boxes is None:
|
||
return []
|
||
return [list(map(float, row)) for row in result.boxes.xyxy.cpu().numpy()]
|
||
|
||
|
||
def main(video_path, weights_path, sample_count=5, device="cpu"):
|
||
model = YOLO(weights_path)
|
||
names = model.names if isinstance(model.names, dict) else dict(enumerate(model.names))
|
||
person_ids = [cid for cid, name in names.items() if name == "person"]
|
||
cap = cv2.VideoCapture(video_path)
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
indices = [int((i + 1) * total / (sample_count + 1)) for i in range(sample_count)]
|
||
direct_count = matched_count = 0
|
||
matched_ious = []
|
||
for index in indices:
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, index)
|
||
ok, frame = cap.read()
|
||
if not ok:
|
||
continue
|
||
h, w = frame.shape[:2]
|
||
precision = 16 if str(device).lower() != "cpu" else None
|
||
direct = boxes(model.predict(frame, imgsz=640, conf=.3, iou=.5,
|
||
classes=person_ids, device=device,
|
||
quantize=precision, verbose=False)[0])
|
||
scale = min(1.0, 2560.0 / max(w, h))
|
||
small = (frame if scale == 1.0 else
|
||
cv2.resize(frame, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA))
|
||
ok, encoded = cv2.imencode(".jpg", small, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
realtime = boxes(model.predict(cv2.imdecode(encoded, cv2.IMREAD_COLOR), imgsz=640,
|
||
conf=.3, iou=.5, classes=person_ids,
|
||
device=device, quantize=precision,
|
||
verbose=False)[0]) if ok else []
|
||
realtime = [[v / scale for v in box] for box in realtime]
|
||
unused = set(range(len(realtime)))
|
||
for box in direct:
|
||
direct_count += 1
|
||
candidates = [(iou(box, realtime[j]), j) for j in unused]
|
||
best = max(candidates, default=(0, None))
|
||
if best[0] >= .8:
|
||
matched_count += 1
|
||
matched_ious.append(best[0])
|
||
unused.remove(best[1])
|
||
print("frame=%d direct_person=%d realtime_person=%d" % (index, len(direct), len(realtime)))
|
||
cap.release()
|
||
recall = matched_count / direct_count if direct_count else 1.0
|
||
mean_iou = sum(matched_ious) / len(matched_ious) if matched_ious else 0.0
|
||
print("matched=%d/%d recall=%.3f mean_iou=%.3f" %
|
||
(matched_count, direct_count, recall, mean_iou))
|
||
return 0 if recall >= .95 and (not matched_ious or min(matched_ious) >= .8) else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
if len(sys.argv) < 3:
|
||
raise SystemExit("usage: validate_yolo_replay.py VIDEO WEIGHTS [SAMPLES] [DEVICE]")
|
||
raise SystemExit(main(sys.argv[1], sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else 5,
|
||
sys.argv[4] if len(sys.argv) > 4 else "cpu"))
|