273 lines
10 KiB
Python
273 lines
10 KiB
Python
|
import os
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from rknnlite.api import RKNNLite
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
# ====================== 配置参数 ======================
|
|||
|
|
MODEL_PATH = "seg.rknn" # 转换好的 RKNN 模型路径
|
|||
|
|
SOURCE_IMG_DIR = "/home/hx/yolo/output_masks" # 输入图像目录
|
|||
|
|
OUTPUT_DIR = "/home/hx/yolo/output_masks_rknn" # 输出目录
|
|||
|
|
ROI_COORDS_FILE = "./roi_coordinates/1_rois.txt" # ROI 文件路径 (x,y,w,h)
|
|||
|
|
TARGET_SIZE = 640 # 模型输入大小
|
|||
|
|
CONF_THRESHOLD = 0.25
|
|||
|
|
IOU_THRESHOLD = 0.45 # 注意：NMS 可能在模型内部完成，也可后处理
|
|||
|
|
DEVICE = RKNNLite.NPU_CORE_0
|
|||
|
|
SAVE_TXT = True
|
|||
|
|
SAVE_MASKS = True
|
|||
|
|
VIEW_IMG = False
|
|||
|
|
LINE_WIDTH = 2
|
|||
|
|
|
|||
|
|
# YOLO 输出结构配置（根据你的模型调整）
|
|||
|
|
MASK_PROTO_IDX = 12 # proto 输出索引
|
|||
|
|
MASK_COEFF_IDXS = [3, 7, 11] # mask coefficient 输出
|
|||
|
|
CONF_IDXS = [1, 5, 9] # objectness 或 conf 输出
|
|||
|
|
BOX_IDXS = [0, 4, 8] # bbox 输出
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_roi_coords(txt_path):
|
|||
|
|
"""加载 ROI 文件，格式: x,y,w,h"""
|
|||
|
|
rois = []
|
|||
|
|
if not os.path.exists(txt_path):
|
|||
|
|
raise FileNotFoundError(f"❌ ROI 文件未找到: {txt_path}")
|
|||
|
|
with open(txt_path, 'r') as f:
|
|||
|
|
for line in f:
|
|||
|
|
line = line.strip()
|
|||
|
|
if line and not line.startswith('#'):
|
|||
|
|
try:
|
|||
|
|
x, y, w, h = map(int, line.split(','))
|
|||
|
|
rois.append((x, y, w, h))
|
|||
|
|
print(f"📌 加载 ROI: (x={x}, y={y}, w={w}, h={h})")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"⚠️ 无法解析 ROI 行: '{line}' | 错误: {e}")
|
|||
|
|
return rois
|
|||
|
|
|
|||
|
|
|
|||
|
|
def sigmoid(x):
|
|||
|
|
return 1 / (1 + np.exp(-x))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def letterbox_resize(image, size, bg_color=114):
|
|||
|
|
"""保持宽高比缩放并填充"""
|
|||
|
|
target_w, target_h = size
|
|||
|
|
h, w, _ = image.shape
|
|||
|
|
scale = min(target_w / w, target_h / h)
|
|||
|
|
new_w, new_h = int(w * scale), int(h * scale)
|
|||
|
|
resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
|||
|
|
canvas = np.full((target_h, target_w, 3), bg_color, dtype=np.uint8)
|
|||
|
|
dx = (target_w - new_w) // 2
|
|||
|
|
dy = (target_h - new_h) // 2
|
|||
|
|
canvas[dy:dy+new_h, dx:dx+new_w] = resized
|
|||
|
|
return canvas, scale, dx, dy
|
|||
|
|
|
|||
|
|
|
|||
|
|
def non_max_suppression(boxes, scores, iou_thresh=0.45):
|
|||
|
|
"""简单 NMS 实现"""
|
|||
|
|
if len(boxes) == 0:
|
|||
|
|
return []
|
|||
|
|
boxes = np.array(boxes)
|
|||
|
|
scores = np.array(scores)
|
|||
|
|
order = scores.argsort()[::-1]
|
|||
|
|
keep = []
|
|||
|
|
while len(order) > 0:
|
|||
|
|
i = order[0]
|
|||
|
|
keep.append(i)
|
|||
|
|
xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
|
|||
|
|
yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
|
|||
|
|
xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
|
|||
|
|
yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
|
|||
|
|
w = np.maximum(0.0, xx2 - xx1 + 1)
|
|||
|
|
h = np.maximum(0.0, yy2 - yy1 + 1)
|
|||
|
|
inter = w * h
|
|||
|
|
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
|||
|
|
inds = np.where(ovr <= iou_thresh)[0]
|
|||
|
|
order = order[inds + 1]
|
|||
|
|
return keep
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run_rknn_inference_with_roi():
|
|||
|
|
# 加载模型
|
|||
|
|
rknn = RKNNLite(verbose=True)
|
|||
|
|
ret = rknn.load_rknn(MODEL_PATH)
|
|||
|
|
if ret != 0:
|
|||
|
|
print("❌ 加载 RKNN 模型失败")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
ret = rknn.init_runtime(core_mask=DEVICE)
|
|||
|
|
if ret != 0:
|
|||
|
|
print("❌ 初始化 NPU 运行时失败")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
print(f"✅ 模型 {MODEL_PATH} 加载成功")
|
|||
|
|
|
|||
|
|
# 创建输出目录
|
|||
|
|
output_dir = Path(OUTPUT_DIR)
|
|||
|
|
txt_dir = output_dir / "labels"
|
|||
|
|
mask_dir = output_dir / "masks"
|
|||
|
|
vis_dir = output_dir / "visualize"
|
|||
|
|
for d in [output_dir, txt_dir, mask_dir, vis_dir]:
|
|||
|
|
d.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 加载 ROI
|
|||
|
|
rois = load_roi_coords(ROI_COORDS_FILE)
|
|||
|
|
if len(rois) == 0:
|
|||
|
|
print("❌ 没有有效 ROI，退出。")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 获取图像列表
|
|||
|
|
img_files = list(Path(SOURCE_IMG_DIR).glob("*.jpg")) + \
|
|||
|
|
list(Path(SOURCE_IMG_DIR).glob("*.png"))
|
|||
|
|
|
|||
|
|
for img_path in img_files:
|
|||
|
|
print(f"\n🔍 处理图像: {img_path.name}")
|
|||
|
|
orig_img = cv2.imread(str(img_path))
|
|||
|
|
if orig_img is None:
|
|||
|
|
print(f"❌ 无法读取图像: {img_path}")
|
|||
|
|
continue
|
|||
|
|
h_orig, w_orig = orig_img.shape[:2]
|
|||
|
|
|
|||
|
|
full_vis_img = orig_img.copy()
|
|||
|
|
all_segments = [] # 存储所有归一化后的多边形点
|
|||
|
|
|
|||
|
|
for roi_idx, (x, y, w, h) in enumerate(rois):
|
|||
|
|
if x < 0 or y < 0 or x + w > w_orig or y + h > h_orig:
|
|||
|
|
print(f"⚠️ ROI 越界，跳过: ({x},{y},{w},{h})")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 提取并预处理 ROI
|
|||
|
|
roi_img = orig_img[y:y+h, x:x+w]
|
|||
|
|
if roi_img.size == 0:
|
|||
|
|
print(f"⚠️ 空 ROI: {roi_idx}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
preprocessed, scale, dx, dy = letterbox_resize(roi_img, (TARGET_SIZE, TARGET_SIZE))
|
|||
|
|
infer_input = preprocessed[..., ::-1].astype(np.float32) # BGR -> RGB
|
|||
|
|
infer_input = np.expand_dims(infer_input, axis=0)
|
|||
|
|
|
|||
|
|
# 推理
|
|||
|
|
outputs = rknn.inference(inputs=[infer_input])
|
|||
|
|
|
|||
|
|
# 解析输出
|
|||
|
|
proto = outputs[MASK_PROTO_IDX][0] # [32, 160, 160]
|
|||
|
|
mask_coeffs_list = [outputs[i] for i in MASK_COEFF_IDXS] # list of [1, C, H, W]
|
|||
|
|
conf_list = [outputs[i] for i in CONF_IDXS]
|
|||
|
|
box_list = [outputs[i] for i in BOX_IDXS]
|
|||
|
|
|
|||
|
|
# 合并所有尺度的检测结果
|
|||
|
|
candidates = []
|
|||
|
|
for s_idx in range(len(conf_list)):
|
|||
|
|
conf_map = conf_list[s_idx].flatten()
|
|||
|
|
box_map = box_list[s_idx].reshape(4, -1).T
|
|||
|
|
coeff_map = mask_coeffs_list[s_idx].reshape(-1, conf_map.shape[0])
|
|||
|
|
for j in range(conf_map.shape[0]):
|
|||
|
|
if conf_map[j] > CONF_THRESHOLD:
|
|||
|
|
candidates.append({
|
|||
|
|
'conf': conf_map[j],
|
|||
|
|
'box': box_map[j], # 相对 640 坐标
|
|||
|
|
'coeff': coeff_map[:, j]
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
if not candidates:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 按置信度排序
|
|||
|
|
candidates.sort(key=lambda x: x['conf'], reverse=True)
|
|||
|
|
top_dets = candidates[:100] # 取前100做NMS
|
|||
|
|
|
|||
|
|
boxes = np.array([d['box'] for d in top_dets])
|
|||
|
|
scores = np.array([d['conf'] for d in top_dets])
|
|||
|
|
coeffs = np.array([d['coeff'] for d in top_dets])
|
|||
|
|
|
|||
|
|
# 反算 bounding box 到 ROI 像素空间
|
|||
|
|
# boxes 是 cx,cy,w,h 归一化到 [0,1]？
|
|||
|
|
# 根据你的模型输出结构调整下面逻辑（假设是 cx,cy,w,h in [0,1]）
|
|||
|
|
cx = boxes[:, 0] * TARGET_SIZE
|
|||
|
|
cy = boxes[:, 1] * TARGET_SIZE
|
|||
|
|
bw = boxes[:, 2] * TARGET_SIZE
|
|||
|
|
bh = boxes[:, 3] * TARGET_SIZE
|
|||
|
|
x1 = cx - bw / 2
|
|||
|
|
y1 = cy - bh / 2
|
|||
|
|
x2 = cx + bw / 2
|
|||
|
|
y2 = cy + bh / 2
|
|||
|
|
|
|||
|
|
# 映射回 letterbox 内的实际区域
|
|||
|
|
x1 = (x1 - dx) / scale
|
|||
|
|
y1 = (y1 - dy) / scale
|
|||
|
|
x2 = (x2 - dx) / scale
|
|||
|
|
y2 = (y2 - dy) / scale
|
|||
|
|
|
|||
|
|
# clip to ROI
|
|||
|
|
x1 = np.clip(x1, 0, w)
|
|||
|
|
y1 = np.clip(y1, 0, h)
|
|||
|
|
x2 = np.clip(x2, 0, w)
|
|||
|
|
y2 = np.clip(y2, 0, h)
|
|||
|
|
|
|||
|
|
# NMS
|
|||
|
|
areas = (x2 - x1) * (y2 - y1)
|
|||
|
|
keep_indices = non_max_suppression(np.stack([x1, y1, x2, y2], axis=1), scores, IOU_THRESHOLD)
|
|||
|
|
|
|||
|
|
for i in keep_indices:
|
|||
|
|
det = top_dets[i]
|
|||
|
|
conf = det['conf']
|
|||
|
|
mask_coeff = det['coeff'] # [32]
|
|||
|
|
|
|||
|
|
# 生成 mask
|
|||
|
|
mask_flat = sigmoid(np.matmul(mask_coeff, proto.reshape(32, -1)))
|
|||
|
|
mask = mask_flat.reshape(160, 160)
|
|||
|
|
mask_resized = cv2.resize(mask, (TARGET_SIZE, TARGET_SIZE), interpolation=cv2.INTER_LINEAR)
|
|||
|
|
|
|||
|
|
# 反变换到 ROI 原始尺寸
|
|||
|
|
mask_roi = cv2.resize(mask_resized, (w, h), interpolation=cv2.INTER_LINEAR)
|
|||
|
|
mask_bin = (mask_roi > 0.5).astype(np.uint8)
|
|||
|
|
|
|||
|
|
# 找轮廓
|
|||
|
|
contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|||
|
|
if len(contours) == 0:
|
|||
|
|
continue
|
|||
|
|
largest_contour = max(contours, key=cv2.contourArea)
|
|||
|
|
segment = largest_contour.squeeze().astype(float)
|
|||
|
|
|
|||
|
|
# 映射回原图坐标
|
|||
|
|
segment[:, 0] += x
|
|||
|
|
segment[:, 1] += y
|
|||
|
|
# 归一化到 [0,1]
|
|||
|
|
segment[:, 0] /= w_orig
|
|||
|
|
segment[:, 1] /= h_orig
|
|||
|
|
all_segments.append((0, segment, conf)) # cls_id=0
|
|||
|
|
|
|||
|
|
# 绘制可视化
|
|||
|
|
cv2.drawContours(full_vis_img[y:y+h, x:x+w], [largest_contour], -1, (0, 255, 0), LINE_WIDTH)
|
|||
|
|
cx_int = int((x1[i] + x2[i]) / 2) + x
|
|||
|
|
cy_int = int(y1[i]) + y - 10
|
|||
|
|
cv2.putText(full_vis_img, f'0 {conf:.2f}', (cx_int, cy_int),
|
|||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 1)
|
|||
|
|
|
|||
|
|
# 保存 ROI mask（可选）
|
|||
|
|
if SAVE_MASKS:
|
|||
|
|
mask_canvas = np.zeros((h_orig, w_orig), dtype=np.uint8)
|
|||
|
|
combined = np.zeros_like(mask_bin)
|
|||
|
|
for i in keep_indices:
|
|||
|
|
mask_flat = sigmoid(np.matmul(top_dets[i]['coeff'], proto.reshape(32, -1)))
|
|||
|
|
mask = cv2.resize(mask_flat.reshape(160, 160), (w, h)) > 0.5
|
|||
|
|
combined |= mask
|
|||
|
|
mask_canvas[y:y+h, x:x+w] = (combined * 255).astype(np.uint8)
|
|||
|
|
cv2.imwrite(str(mask_dir / f"{img_path.stem}_roi{roi_idx}.png"), mask_canvas)
|
|||
|
|
|
|||
|
|
# 保存最终可视化
|
|||
|
|
cv2.imwrite(str(vis_dir / f"vis_{img_path.name}"), full_vis_img)
|
|||
|
|
|
|||
|
|
# 保存 TXT 标签
|
|||
|
|
if SAVE_TXT and all_segments:
|
|||
|
|
with open(txt_dir / f"{img_path.stem}.txt", 'w') as f:
|
|||
|
|
for cls_id, seg, conf in all_segments:
|
|||
|
|
seg_flat = seg.flatten()
|
|||
|
|
f.write(f"{cls_id} {' '.join(f'{x:.6f}' for x in seg_flat)}\n")
|
|||
|
|
|
|||
|
|
print(f"✅ 已保存结果: {vis_dir / f'vis_{img_path.name}'}")
|
|||
|
|
|
|||
|
|
rknn.release()
|
|||
|
|
print(f"\n🎉 全部完成！输出位于: {OUTPUT_DIR}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
run_rknn_inference_with_roi()
|