71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
# ------------------ 手动定义分类标签 ------------------
|
|
labels = ["noready", "ready"] # 0: 盖板不到位, 1: 盖板到位
|
|
|
|
# ------------------ 配置 ------------------
|
|
IMG_SIZE = (640, 640)
|
|
|
|
|
|
def preprocess_cls(bgr_image_corrected, target_size=IMG_SIZE):
|
|
"""
|
|
对已经校正方向（翻转后）的 BGR 图像进行预处理：
|
|
- 转 RGB
|
|
- Resize
|
|
- 添加 batch 维度
|
|
注意：此函数不再负责翻转！翻转应在调用前完成。
|
|
"""
|
|
rgb = cv2.cvtColor(bgr_image_corrected, cv2.COLOR_BGR2RGB)
|
|
resized = cv2.resize(rgb, target_size)
|
|
input_tensor = np.expand_dims(resized, axis=0).astype(np.uint8)
|
|
return input_tensor
|
|
|
|
|
|
def get_top1_class_str_from_output(outputs):
|
|
if outputs is None or len(outputs) == 0:
|
|
return "Error: No output"
|
|
logits = outputs[0].flatten()
|
|
top1_idx = int(np.argmax(logits))
|
|
if 0 <= top1_idx < len(labels):
|
|
return labels[top1_idx]
|
|
else:
|
|
return "Unknown"
|
|
|
|
|
|
def myFunc(rknn_lite, IMG):
|
|
"""
|
|
统一推理接口（分类任务）
|
|
- 输入原始 BGR 图像（可能倒置）
|
|
- 内部先上下翻转（校正方向）
|
|
- 模型在翻转后的图像上推理
|
|
- 返回带结果的翻转后图像（即“正”的图像）
|
|
"""
|
|
# 1. 上下翻转原始图像（校正摄像头倒装）
|
|
corrected_img = cv2.flip(IMG, 0) # 现在这是“正”的图像
|
|
|
|
# 2. 预处理（使用校正后的图像）
|
|
input_data = preprocess_cls(corrected_img, target_size=IMG_SIZE)
|
|
|
|
# 3. 推理
|
|
outputs = rknn_lite.inference(inputs=[input_data])
|
|
|
|
# 4. 解析结果
|
|
class_name = get_top1_class_str_from_output(outputs)
|
|
|
|
# 5. 在校正后的图像上绘制结果
|
|
vis_img = corrected_img.copy()
|
|
text = f"Class: {class_name}"
|
|
cv2.putText(
|
|
vis_img,
|
|
text,
|
|
(10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.0,
|
|
(0, 255, 0),
|
|
2,
|
|
cv2.LINE_AA
|
|
)
|
|
|
|
return vis_img # 返回的是上下翻转后的 BGR 图像（即视觉上“正”的图）
|