76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
|
import os
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
from rknnlite.api import RKNNLite
|
||
|
|
|
||
|
|
objectThresh = 0.5 # 置信度阈值，可调整
|
||
|
|
|
||
|
|
def letterbox_resize(image, size, bg_color=114):
|
||
|
|
target_width, target_height = size
|
||
|
|
h, w, _ = image.shape
|
||
|
|
scale = min(target_width / w, target_height / h)
|
||
|
|
new_w, new_h = int(w * scale), int(h * scale)
|
||
|
|
resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||
|
|
canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
|
||
|
|
offset_x, offset_y = (target_width - new_w) // 2, (target_height - new_h) // 2
|
||
|
|
canvas[offset_y:offset_y+new_h, offset_x:offset_x+new_w] = resized
|
||
|
|
return canvas, scale, offset_x, offset_y
|
||
|
|
|
||
|
|
def sigmoid(x):
|
||
|
|
return 1 / (1 + np.exp(-x))
|
||
|
|
|
||
|
|
def detect_top2_single_masks(model_path, image_path):
|
||
|
|
img = cv2.imread(image_path)
|
||
|
|
if img is None:
|
||
|
|
print(f"❌ 无法读取图像 {image_path}")
|
||
|
|
return []
|
||
|
|
|
||
|
|
img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640))
|
||
|
|
infer_img = img_resized[..., ::-1] # BGR->RGB
|
||
|
|
infer_img = np.expand_dims(infer_img, 0)
|
||
|
|
|
||
|
|
rknn = RKNNLite(verbose=True)
|
||
|
|
rknn.load_rknn(model_path)
|
||
|
|
rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||
|
|
outputs = rknn.inference([infer_img])
|
||
|
|
rknn.release()
|
||
|
|
|
||
|
|
mask_coeffs_list = [outputs[3], outputs[7], outputs[11]] # mask coefficients
|
||
|
|
conf_list = [outputs[1], outputs[5], outputs[9]] # object confidence
|
||
|
|
proto = outputs[12][0] # proto [32,160,160]
|
||
|
|
|
||
|
|
# 遍历所有候选，生成(置信度, 尺度索引, 目标索引)
|
||
|
|
candidates = []
|
||
|
|
for scale_idx, conf_map in enumerate(conf_list):
|
||
|
|
conf_map_flat = conf_map.flatten()
|
||
|
|
for idx, conf in enumerate(conf_map_flat):
|
||
|
|
if conf > objectThresh:
|
||
|
|
candidates.append((conf, scale_idx, idx))
|
||
|
|
|
||
|
|
if not candidates:
|
||
|
|
print(f"⚠️ 未检测到目标，置信度低于 {objectThresh}")
|
||
|
|
return []
|
||
|
|
|
||
|
|
# 按置信度排序，取前两个
|
||
|
|
candidates.sort(key=lambda x: x[0], reverse=True)
|
||
|
|
top2 = candidates[:2]
|
||
|
|
|
||
|
|
masks = []
|
||
|
|
for i, (conf, scale_idx, idx) in enumerate(top2):
|
||
|
|
coeff = mask_coeffs_list[scale_idx].reshape(mask_coeffs_list[scale_idx].shape[1], -1)
|
||
|
|
mask_flat = np.matmul(coeff[:, idx], proto.reshape(proto.shape[0], -1))
|
||
|
|
mask = sigmoid(mask_flat).reshape(proto.shape[1], proto.shape[2])
|
||
|
|
mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]))
|
||
|
|
mask_bin = (mask_resized > 0.5).astype(np.uint8) * 255
|
||
|
|
masks.append(mask_bin)
|
||
|
|
cv2.imwrite(f"mask_single_{i+1}.png", mask_bin)
|
||
|
|
print(f"✅ mask {i+1} 已保存: mask_single_{i+1}.png")
|
||
|
|
|
||
|
|
return masks
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
model_path = "seg.rknn"
|
||
|
|
image_path = "2.jpg"
|
||
|
|
detect_top2_single_masks(model_path, image_path)
|
||
|
|
|