131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
|
|
TARGET_SIZE = 640 # 模型输入尺寸
|
|
|
|
# --------------------
|
|
# 全局 ROI 定义
|
|
# --------------------
|
|
ROIS = [
|
|
(859,810,696,328), # (x, y, w, h)
|
|
]
|
|
|
|
# --------------------
|
|
# 剔除相邻 x 差距过大的离散点
|
|
# --------------------
|
|
def filter_outliers_by_x(pts, x_thresh=30):
|
|
if len(pts) < 2:
|
|
return pts
|
|
pts_sorted = pts[np.argsort(pts[:, 1])] # 按 y 排序
|
|
clean_pts = [pts_sorted[0]]
|
|
for i in range(1, len(pts_sorted)):
|
|
if abs(pts_sorted[i, 0] - pts_sorted[i-1, 0]) < x_thresh:
|
|
clean_pts.append(pts_sorted[i])
|
|
return np.array(clean_pts, dtype=np.int32)
|
|
|
|
# --------------------
|
|
# 根据角点分布，选取左右边缘角点
|
|
# --------------------
|
|
def select_edge_corners(corners, w, left_ratio=0.2, right_ratio=0.2, y_var_thresh=5, x_var_thresh=30):
|
|
if corners is None:
|
|
return [], []
|
|
|
|
corners = np.int32(corners).reshape(-1, 2)
|
|
x_min, x_max = 0, w
|
|
|
|
left_thresh = x_min + int(w * left_ratio)
|
|
right_thresh = x_max - int(w * right_ratio)
|
|
|
|
# 左右候选角点
|
|
left_candidates = corners[corners[:, 0] <= left_thresh]
|
|
right_candidates = corners[corners[:, 0] >= right_thresh]
|
|
|
|
# --------------------
|
|
# 进一步按 y 变化筛选
|
|
# --------------------
|
|
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)
|
|
|
|
# --------------------
|
|
# 再剔除相邻 x 值差距过大的离散点
|
|
# --------------------
|
|
left_final = filter_outliers_by_x(left_final, x_var_thresh)
|
|
right_final = filter_outliers_by_x(right_final, x_var_thresh)
|
|
|
|
return left_final, right_final
|
|
|
|
# --------------------
|
|
# 推理 ROI 并可视化 mask + 边缘角点
|
|
# --------------------
|
|
def infer_mask_with_selected_corners(image_path, model_path, output_dir="./output"):
|
|
model = YOLO(model_path)
|
|
|
|
image_path = Path(image_path)
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
orig_img = cv2.imread(str(image_path))
|
|
overlay_img = orig_img.copy()
|
|
|
|
for idx, (x, y, w, h) in enumerate(ROIS):
|
|
roi_img = orig_img[y:y+h, x:x+w]
|
|
resized_img = cv2.resize(roi_img, (TARGET_SIZE, TARGET_SIZE))
|
|
|
|
# 模型推理
|
|
results = model(source=resized_img, imgsz=TARGET_SIZE, verbose=False)
|
|
result = results[0]
|
|
|
|
# 可视化 mask
|
|
if result.masks is not None and len(result.masks.data) > 0:
|
|
mask = result.masks.data[0].cpu().numpy()
|
|
mask_bin = (mask > 0.5).astype(np.uint8)
|
|
mask_bin = cv2.resize(mask_bin, (w, h), interpolation=cv2.INTER_NEAREST)
|
|
|
|
# 绿色 mask 覆盖
|
|
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)
|
|
|
|
# 可视化
|
|
for cx, cy in left_pts:
|
|
cv2.circle(overlay_img[y:y+h, x:x+w], (cx, cy), 6, (0, 0, 255), -1) # 左边红色
|
|
for cx, cy in right_pts:
|
|
cv2.circle(overlay_img[y:y+h, x:x+w], (cx, cy), 6, (255, 0, 0), -1) # 右边蓝色
|
|
|
|
# 保存结果
|
|
save_path = output_dir / f"mask_edge_corners_{image_path.name}"
|
|
cv2.imwrite(str(save_path), overlay_img)
|
|
print(f"✅ 保存结果: {save_path}")
|
|
|
|
return overlay_img
|
|
|
|
|
|
# ===================== 使用示例 =====================
|
|
if __name__ == "__main__":
|
|
IMAGE_PATH = "../test_image/1.jpg"
|
|
MODEL_PATH = "best.pt"
|
|
|
|
infer_mask_with_selected_corners(IMAGE_PATH, MODEL_PATH)
|