276 lines
7.2 KiB
Python
276 lines
7.2 KiB
Python
import os
|
|
import cv2
|
|
import time
|
|
import numpy as np
|
|
from datetime import datetime
|
|
from collections import deque
|
|
from rknnlite.api import RKNNLite
|
|
|
|
# =====================================================
|
|
# 稳定判决器
|
|
# =====================================================
|
|
class StableClassJudge:
|
|
"""
|
|
连续 N 帧稳定判决：
|
|
- class0 / class1 连续 N 帧 -> 输出
|
|
- class2 -> 清空计数
|
|
"""
|
|
|
|
def __init__(self, stable_frames=3, ignore_class=2):
|
|
self.stable_frames = stable_frames
|
|
self.ignore_class = ignore_class
|
|
self.buffer = deque(maxlen=stable_frames)
|
|
|
|
def reset(self):
|
|
self.buffer.clear()
|
|
|
|
def update(self, class_id):
|
|
if class_id == self.ignore_class:
|
|
self.reset()
|
|
return None
|
|
|
|
self.buffer.append(class_id)
|
|
|
|
if len(self.buffer) < self.stable_frames:
|
|
return None
|
|
|
|
if len(set(self.buffer)) == 1:
|
|
stable = self.buffer[0]
|
|
self.reset()
|
|
return stable
|
|
|
|
return None
|
|
|
|
|
|
# =====================================================
|
|
# 类别定义
|
|
# =====================================================
|
|
CLASS_NAMES = {
|
|
0: "未插入振捣棒",
|
|
1: "安全插入振捣棒",
|
|
2: "有遮挡"
|
|
}
|
|
|
|
|
|
# =====================================================
|
|
# RKNN 全局实例
|
|
# =====================================================
|
|
_global_rknn = None
|
|
|
|
|
|
def init_rknn_model(model_path):
|
|
global _global_rknn
|
|
if _global_rknn is not None:
|
|
return _global_rknn
|
|
|
|
rknn = RKNNLite(verbose=False)
|
|
|
|
ret = rknn.load_rknn(model_path)
|
|
if ret != 0:
|
|
raise RuntimeError(f"Load RKNN failed: {ret}")
|
|
|
|
ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
|
if ret != 0:
|
|
raise RuntimeError(f"Init runtime failed: {ret}")
|
|
|
|
_global_rknn = rknn
|
|
print(f"[INFO] RKNN 模型加载成功: {model_path}")
|
|
return rknn
|
|
|
|
|
|
# =====================================================
|
|
# 图像预处理
|
|
# =====================================================
|
|
def letterbox(image, new_size=640, color=(114, 114, 114)):
|
|
h, w = image.shape[:2]
|
|
scale = min(new_size / h, new_size / w)
|
|
nh, nw = int(h * scale), int(w * scale)
|
|
|
|
resized = cv2.resize(image, (nw, nh))
|
|
canvas = np.full((new_size, new_size, 3), color, dtype=np.uint8)
|
|
|
|
top = (new_size - nh) // 2
|
|
left = (new_size - nw) // 2
|
|
canvas[top:top + nh, left:left + nw] = resized
|
|
return canvas
|
|
|
|
|
|
def resize_stretch(image, size=640):
|
|
return cv2.resize(image, (size, size))
|
|
|
|
|
|
def preprocess_image_for_rknn(
|
|
img,
|
|
size=640,
|
|
resize_mode="stretch",
|
|
to_rgb=True,
|
|
normalize=False,
|
|
layout="NHWC"
|
|
):
|
|
if resize_mode == "letterbox":
|
|
img = letterbox(img, size)
|
|
else:
|
|
img = resize_stretch(img, size)
|
|
|
|
if to_rgb:
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
img = img.astype(np.float32)
|
|
|
|
if normalize:
|
|
img /= 255.0
|
|
|
|
if layout == "NHWC":
|
|
img = np.expand_dims(img, axis=0)
|
|
else:
|
|
img = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
|
|
|
|
return np.ascontiguousarray(img)
|
|
|
|
|
|
# =====================================================
|
|
# RKNN 单次推理
|
|
# =====================================================
|
|
def rknn_classify_preprocessed(input_tensor, model_path):
|
|
rknn = init_rknn_model(model_path)
|
|
outs = rknn.inference([input_tensor])
|
|
probs = outs[0].reshape(-1).astype(float)
|
|
class_id = int(np.argmax(probs))
|
|
return class_id, probs
|
|
|
|
|
|
# =====================================================
|
|
# ROI 处理
|
|
# =====================================================
|
|
def load_single_roi(txt_path):
|
|
if not os.path.exists(txt_path):
|
|
raise RuntimeError(f"ROI 文件不存在: {txt_path}")
|
|
|
|
with open(txt_path) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
x, y, w, h = map(int, line.split(","))
|
|
return (x, y, w, h)
|
|
|
|
raise RuntimeError("ROI 文件为空")
|
|
|
|
|
|
def crop_and_return_roi(img, roi):
|
|
x, y, w, h = roi
|
|
H, W = img.shape[:2]
|
|
|
|
if x < 0 or y < 0 or x + w > W or y + h > H:
|
|
raise RuntimeError(f"ROI 超出图像范围: {roi}")
|
|
|
|
return img[y:y + h, x:x + w]
|
|
|
|
|
|
# =====================================================
|
|
# 单帧分类
|
|
# =====================================================
|
|
def classify_single_image(frame, model_path, roi_file):
|
|
roi = load_single_roi(roi_file)
|
|
roi_img = crop_and_return_roi(frame, roi)
|
|
|
|
input_tensor = preprocess_image_for_rknn(
|
|
roi_img,
|
|
size=640,
|
|
resize_mode="stretch",
|
|
to_rgb=True,
|
|
normalize=False,
|
|
layout="NHWC"
|
|
)
|
|
|
|
class_id, probs = rknn_classify_preprocessed(input_tensor, model_path)
|
|
|
|
return {
|
|
"class_id": class_id,
|
|
"class": CLASS_NAMES[class_id],
|
|
"score": round(float(probs[class_id]), 4),
|
|
"raw": probs.tolist()
|
|
}
|
|
|
|
|
|
# =====================================================
|
|
# RTSP 推理 + 保存分类结果
|
|
# =====================================================
|
|
def run_rtsp_classification_and_save(
|
|
model_path,
|
|
roi_file,
|
|
rtsp_url,
|
|
save_root="clsimg",
|
|
stable_frames=3,
|
|
save_mode="all" # all / stable
|
|
):
|
|
for cid in CLASS_NAMES.keys():
|
|
os.makedirs(os.path.join(save_root, f"class{cid}"), exist_ok=True)
|
|
|
|
cap = cv2.VideoCapture(rtsp_url)
|
|
if not cap.isOpened():
|
|
raise RuntimeError(f"无法打开 RTSP: {rtsp_url}")
|
|
|
|
judge = StableClassJudge(stable_frames=stable_frames, ignore_class=2)
|
|
|
|
print("[INFO] RTSP 推理开始")
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("[WARN] RTSP 读帧失败")
|
|
time.sleep(0.1)
|
|
continue
|
|
|
|
frame = cv2.flip(frame, -1)
|
|
|
|
result = classify_single_image(frame, model_path, roi_file)
|
|
class_id = result["class_id"]
|
|
score = result["score"]
|
|
|
|
print(f"[FRAME] {result['class']} conf={score}")
|
|
|
|
stable = judge.update(class_id)
|
|
|
|
save_flag = False
|
|
save_class = class_id
|
|
|
|
if save_mode == "all":
|
|
save_flag = True
|
|
elif save_mode == "stable" and stable is not None:
|
|
save_flag = True
|
|
save_class = stable
|
|
|
|
if save_flag:
|
|
ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
|
filename = f"{ts}_conf{score:.2f}.jpg"
|
|
save_dir = os.path.join(save_root, f"class{save_class}")
|
|
cv2.imwrite(os.path.join(save_dir, filename), frame)
|
|
print(f"[SAVE] class{save_class}/{filename}")
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
# =====================================================
|
|
# main
|
|
# =====================================================
|
|
if __name__ == "__main__":
|
|
model_path = "zdb_cls.rknn"
|
|
roi_file = "./roi_coordinates/zdb_roi.txt"
|
|
|
|
rtsp_url = "rtsp://admin:XJ123456@192.168.250.60:554/streaming/channels/101"
|
|
|
|
run_rtsp_classification_and_save(
|
|
model_path=model_path,
|
|
roi_file=roi_file,
|
|
rtsp_url=rtsp_url,
|
|
save_root="clsimg",
|
|
stable_frames=3,
|
|
save_mode="all" # 改成 "stable" 只存稳定结果
|
|
)
|
|
|