244 lines
8.0 KiB
Python
244 lines
8.0 KiB
Python
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from pathlib import Path
|
|||
|
|
from rknnlite.api import RKNNLite
|
|||
|
|
|
|||
|
|
# --------------------
|
|||
|
|
# 配置参数
|
|||
|
|
# --------------------
|
|||
|
|
IMAGE_PATH = "./11.jpg"
|
|||
|
|
MODEL_PATH = "segr.rknn"
|
|||
|
|
OUTPUT_PATH = "./single_result.jpg"
|
|||
|
|
|
|||
|
|
TARGET_SIZE = 640
|
|||
|
|
|
|||
|
|
# 像素到实际尺寸换算比例
|
|||
|
|
def calculate_pixel_to_real_ratio(real_length_mm, pixel_length):
|
|||
|
|
if pixel_length == 0:
|
|||
|
|
raise ValueError("像素长度不能为0")
|
|||
|
|
return real_length_mm / pixel_length
|
|||
|
|
|
|||
|
|
PIXEL_TO_REAL_RATIO = 1.0
|
|||
|
|
real_length_mm = 100 # 实际长度（毫米）
|
|||
|
|
pixel_length = 200 # 对应像素长度
|
|||
|
|
try:
|
|||
|
|
PIXEL_TO_REAL_RATIO = calculate_pixel_to_real_ratio(real_length_mm, pixel_length)
|
|||
|
|
print(f"换算比例已设定: {PIXEL_TO_REAL_RATIO:.4f} mm/像素")
|
|||
|
|
except ValueError as e:
|
|||
|
|
print(e)
|
|||
|
|
|
|||
|
|
# 全局 ROI 定义：(x, y, w, h)
|
|||
|
|
ROIS = [
|
|||
|
|
(859, 810, 696, 328),
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# --------------------
|
|||
|
|
# RKNN 单例管理
|
|||
|
|
# --------------------
|
|||
|
|
_rknn_instance = None
|
|||
|
|
|
|||
|
|
def init_rknn(model_path):
|
|||
|
|
global _rknn_instance
|
|||
|
|
if _rknn_instance is not None:
|
|||
|
|
return _rknn_instance
|
|||
|
|
|
|||
|
|
_rknn_instance = RKNNLite(verbose=False)
|
|||
|
|
ret = _rknn_instance.load_rknn(model_path)
|
|||
|
|
if ret != 0:
|
|||
|
|
print(f"[ERROR] 加载 RKNN 模型失败: {ret}")
|
|||
|
|
_rknn_instance = None
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
ret = _rknn_instance.init_runtime(core_mask=RKNNLite.NPU_CORE_1)
|
|||
|
|
if ret != 0:
|
|||
|
|
print(f"[ERROR] 初始化 NPU 失败: {ret}")
|
|||
|
|
_rknn_instance.release()
|
|||
|
|
_rknn_instance = None
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
print("[✅] RKNN 分割模型加载成功")
|
|||
|
|
return _rknn_instance
|
|||
|
|
|
|||
|
|
def release_rknn():
|
|||
|
|
global _rknn_instance
|
|||
|
|
if _rknn_instance:
|
|||
|
|
_rknn_instance.release()
|
|||
|
|
_rknn_instance = None
|
|||
|
|
print("[INFO] RKNN 模型已释放")
|
|||
|
|
|
|||
|
|
# --------------------
|
|||
|
|
# 工具函数：letterbox resize
|
|||
|
|
# --------------------
|
|||
|
|
def letterbox_resize(image, size, bg_color=114):
|
|||
|
|
target_w, target_h = size
|
|||
|
|
h, w = image.shape[:2]
|
|||
|
|
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 select_edge_corners(corners, w, left_ratio=0.2, right_ratio=0.2, y_var_thresh=5):
|
|||
|
|
if corners is None:
|
|||
|
|
return [], []
|
|||
|
|
corners = np.int32(corners).reshape(-1, 2)
|
|||
|
|
left_thresh = int(w * left_ratio)
|
|||
|
|
right_thresh = w - int(w * right_ratio)
|
|||
|
|
|
|||
|
|
left_candidates = corners[corners[:, 0] <= left_thresh]
|
|||
|
|
right_candidates = corners[corners[:, 0] >= right_thresh]
|
|||
|
|
|
|||
|
|
def filter_by_y_variation(pts):
|
|||
|
|
if len(pts) < 2:
|
|||
|
|
return pts
|
|||
|
|
pts_sorted = pts[np.argsort(pts[:, 1])]
|
|||
|
|
diffs = np.abs(np.diff(pts_sorted[:, 1]))
|
|||
|
|
keep_idx = np.where(diffs > y_var_thresh)[0]
|
|||
|
|
selected = [pts_sorted[i] for i in keep_idx] + [pts_sorted[i + 1] for i in keep_idx]
|
|||
|
|
return np.array(selected) if len(selected) > 0 else pts_sorted
|
|||
|
|
|
|||
|
|
left_final = filter_by_y_variation(left_candidates)
|
|||
|
|
right_final = filter_by_y_variation(right_candidates)
|
|||
|
|
return left_final, right_final
|
|||
|
|
|
|||
|
|
def fit_line_with_outlier_removal(pts, dist_thresh=10):
|
|||
|
|
if pts is None or len(pts) < 2:
|
|||
|
|
return None, pts
|
|||
|
|
pts = np.array(pts)
|
|||
|
|
x, y = pts[:, 0], pts[:, 1]
|
|||
|
|
m, b = np.polyfit(y, x, 1) # x = m*y + b
|
|||
|
|
x_fit = m * y + b
|
|||
|
|
dists = np.abs(x - x_fit)
|
|||
|
|
mask = dists < dist_thresh
|
|||
|
|
if mask.sum() < 2:
|
|||
|
|
return (m, b), pts
|
|||
|
|
m, b = np.polyfit(y[mask], x[mask], 1)
|
|||
|
|
inliers = np.stack([x[mask], y[mask]], axis=1)
|
|||
|
|
return (m, b), inliers
|
|||
|
|
|
|||
|
|
# --------------------
|
|||
|
|
# RKNN 推理生成 mask（增加去除小区域）
|
|||
|
|
# --------------------
|
|||
|
|
def get_mask_from_rknn(rknn_model, roi_img, target_size=640, min_area=100):
|
|||
|
|
h_orig, w_orig = roi_img.shape[:2]
|
|||
|
|
preprocessed, scale, dx, dy = letterbox_resize(roi_img, (target_size, target_size))
|
|||
|
|
infer_input = preprocessed[np.newaxis, :, :, ::-1].astype(np.float32)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
outputs = rknn_model.inference(inputs=[infer_input])
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[ERROR] RKNN 推理异常: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
proto = outputs[12][0] # (32,160,160)
|
|||
|
|
mask_proto = np.mean(proto, axis=0)
|
|||
|
|
mask_proto = 1 / (1 + np.exp(-mask_proto))
|
|||
|
|
|
|||
|
|
mask_lb = cv2.resize(mask_proto, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
|
|||
|
|
|
|||
|
|
scale = min(target_size / w_orig, target_size / h_orig)
|
|||
|
|
new_w, new_h = int(w_orig*scale), int(h_orig*scale)
|
|||
|
|
pad_x = (target_size - new_w)//2
|
|||
|
|
pad_y = (target_size - new_h)//2
|
|||
|
|
mask_cropped = mask_lb[pad_y:pad_y+new_h, pad_x:pad_x+new_w]
|
|||
|
|
|
|||
|
|
mask_resized = cv2.resize(mask_cropped, (w_orig, h_orig), interpolation=cv2.INTER_LINEAR)
|
|||
|
|
mask_bin = (mask_resized > 0.5).astype(np.uint8)
|
|||
|
|
|
|||
|
|
# 去除小区域
|
|||
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_bin, connectivity=8)
|
|||
|
|
mask_clean = np.zeros_like(mask_bin)
|
|||
|
|
for i in range(1, num_labels):
|
|||
|
|
area = stats[i, cv2.CC_STAT_AREA]
|
|||
|
|
if area >= min_area:
|
|||
|
|
mask_clean[labels == i] = 1
|
|||
|
|
|
|||
|
|
return mask_clean
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[ERROR] 生成 mask 失败: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# --------------------
|
|||
|
|
# 主函数
|
|||
|
|
# --------------------
|
|||
|
|
def infer_single_image(image_path, model_path, output_path):
|
|||
|
|
orig_img = cv2.imread(str(image_path))
|
|||
|
|
if orig_img is None:
|
|||
|
|
print(f"无法读取图像: {image_path}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
overlay_img = orig_img.copy()
|
|||
|
|
x_diff_pixel = None
|
|||
|
|
|
|||
|
|
rknn = init_rknn(model_path)
|
|||
|
|
if rknn is None:
|
|||
|
|
print("RKNN 初始化失败")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
for idx, (x, y, w, h) in enumerate(ROIS):
|
|||
|
|
roi_img = orig_img[y:y+h, x:x+w]
|
|||
|
|
|
|||
|
|
mask_bin = get_mask_from_rknn(rknn, roi_img, target_size=TARGET_SIZE, min_area=30000)
|
|||
|
|
if mask_bin is None or mask_bin.sum() == 0:
|
|||
|
|
print("未检测到有效 mask")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
color_mask = np.zeros_like(roi_img, dtype=np.uint8)
|
|||
|
|
color_mask[mask_bin == 1] = (0, 255, 0)
|
|||
|
|
overlay_img[y:y+h, x:x+w] = cv2.addWeighted(roi_img, 0.7, color_mask, 0.3, 0)
|
|||
|
|
|
|||
|
|
mask_gray = (mask_bin * 255).astype(np.uint8)
|
|||
|
|
corners = cv2.goodFeaturesToTrack(mask_gray, maxCorners=200, qualityLevel=0.01, minDistance=5)
|
|||
|
|
|
|||
|
|
left_pts, right_pts = select_edge_corners(corners, w)
|
|||
|
|
left_line, _ = fit_line_with_outlier_removal(left_pts)
|
|||
|
|
right_line, _ = fit_line_with_outlier_removal(right_pts)
|
|||
|
|
|
|||
|
|
if left_line and right_line:
|
|||
|
|
y_ref = h * 0.6
|
|||
|
|
m1, b1 = left_line
|
|||
|
|
m2, b2 = right_line
|
|||
|
|
x1 = m1 * y_ref + b1
|
|||
|
|
x2 = m2 * y_ref + b2
|
|||
|
|
x_diff_pixel = abs(x2 - x1)
|
|||
|
|
|
|||
|
|
cv2.line(overlay_img[y:y+h, x:x+w], (0, int(y_ref)), (w, int(y_ref)), (0, 255, 255), 2)
|
|||
|
|
cv2.putText(overlay_img[y:y+h, x:x+w],
|
|||
|
|
f"x_diff={x_diff_pixel:.1f}px",
|
|||
|
|
(10, 30), cv2.FONT_HERSHEY_SIMPLEX,
|
|||
|
|
1, (0, 255, 255), 2)
|
|||
|
|
|
|||
|
|
for (m, b), color in [(left_line, (0, 0, 255)), (right_line, (255, 0, 0))]:
|
|||
|
|
y1, y2 = 0, h
|
|||
|
|
x1_line, x2_line = int(m * y1 + b), int(m * y2 + b)
|
|||
|
|
cv2.line(overlay_img[y:y+h, x:x+w], (x1_line, y1), (x2_line, y2), color, 3)
|
|||
|
|
|
|||
|
|
cv2.imwrite(output_path, overlay_img)
|
|||
|
|
print(f"结果已保存至: {output_path}")
|
|||
|
|
|
|||
|
|
if x_diff_pixel is not None:
|
|||
|
|
x_diff_real = x_diff_pixel * PIXEL_TO_REAL_RATIO
|
|||
|
|
print(f"x差值（像素） = {x_diff_pixel:.2f} px")
|
|||
|
|
print(f"x差值（实际） = {x_diff_real:.2f} mm")
|
|||
|
|
else:
|
|||
|
|
print("未能计算 x 差值")
|
|||
|
|
|
|||
|
|
release_rknn()
|
|||
|
|
return x_diff_pixel
|
|||
|
|
|
|||
|
|
# =====================
|
|||
|
|
# 运行入口
|
|||
|
|
# =====================
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
infer_single_image(IMAGE_PATH, MODEL_PATH, OUTPUT_PATH)
|
|||
|
|
|