212 lines
6.7 KiB
Python
212 lines
6.7 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
import matplotlib.pyplot as plt
|
|
|
|
TARGET_SIZE = 640 # 模型输入尺寸
|
|
IMAGE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/611/train/class0" # 👈 图片文件夹路径
|
|
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/61seg/exp5/weights/best.pt"
|
|
OUTPUT_DIR = "./output"
|
|
|
|
# --------------------
|
|
# 全局 ROI 定义
|
|
# --------------------
|
|
ROIS = [
|
|
(670,623,465,178), # (x, y, w, h)
|
|
]
|
|
|
|
# --------------------
|
|
# 根据角点分布，选取左右边缘角点
|
|
# --------------------
|
|
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)
|
|
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)
|
|
|
|
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
|
|
x2, y2 = x[mask], y[mask]
|
|
|
|
if len(x2) < 2:
|
|
return (m, b), pts
|
|
|
|
# 二次拟合
|
|
m, b = np.polyfit(y2, x2, 1)
|
|
return (m, b), np.stack([x2, y2], axis=1)
|
|
|
|
|
|
# --------------------
|
|
# 推理 ROI 并计算左右直线 x 差
|
|
# --------------------
|
|
def infer_mask_and_get_xdiff(image_path, model, output_dir="./output"):
|
|
image_path = Path(image_path)
|
|
orig_img = cv2.imread(str(image_path))
|
|
overlay_img = orig_img.copy()
|
|
|
|
x_diff_value = None
|
|
intersections = [] # 👈 保存交点 (全图坐标)
|
|
|
|
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]
|
|
|
|
if result.masks is None or len(result.masks.data) == 0:
|
|
print(f"❌ {image_path.name} 未检测到 mask")
|
|
continue
|
|
|
|
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_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 = int(h * 0.6)
|
|
|
|
# ---------- 左右交点 ----------
|
|
m1, b1 = left_line
|
|
m2, b2 = right_line
|
|
|
|
x_left = int(m1 * y_ref + b1)
|
|
x_right = int(m2 * y_ref + b2)
|
|
x_diff_value = abs(x_right - x_left)
|
|
|
|
# ---------- ROI → 原图坐标 ----------
|
|
left_pt_global = (x + x_left, y + y_ref)
|
|
right_pt_global = (x + x_right, y + y_ref)
|
|
|
|
intersections.append({
|
|
"left": left_pt_global,
|
|
"right": right_pt_global,
|
|
"x_diff": x_diff_value
|
|
})
|
|
|
|
# ---------- 画黄色参考线 ----------
|
|
cv2.line(
|
|
overlay_img[y:y+h, x:x+w],
|
|
(0, y_ref), (w, y_ref),
|
|
(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.circle(overlay_img, left_pt_global, 6, (0, 0, 255), -1)
|
|
cv2.circle(overlay_img, right_pt_global, 6, (255, 0, 0), -1)
|
|
|
|
cv2.putText(
|
|
overlay_img,
|
|
f"L{left_pt_global} R{right_pt_global}",
|
|
(x + 10, y + 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.8,
|
|
(0, 255, 255),
|
|
2
|
|
)
|
|
|
|
save_path = Path(output_dir) / f"mask_edge_{image_path.name}"
|
|
cv2.imwrite(str(save_path), overlay_img)
|
|
print(f"✅ 保存结果: {save_path}")
|
|
|
|
return x_diff_value, intersections
|
|
|
|
|
|
|
|
# =====================
|
|
# 批量处理文件夹中的图片
|
|
# =====================
|
|
def process_image_folder(image_dir, model_path):
|
|
model = YOLO(model_path)
|
|
output_dir = OUTPUT_DIR
|
|
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
xdiff_results = {}
|
|
|
|
for img_name in sorted(os.listdir(image_dir)):
|
|
if not img_name.lower().endswith((".jpg", ".png", ".jpeg")):
|
|
continue
|
|
img_path = os.path.join(image_dir, img_name)
|
|
x_diff = infer_mask_and_get_xdiff(img_path, model, output_dir)
|
|
if x_diff is not None:
|
|
xdiff_results[img_name] = x_diff
|
|
|
|
# 结果展示
|
|
if xdiff_results:
|
|
plt.figure(figsize=(8, 4))
|
|
plt.plot(list(xdiff_results.keys()), list(xdiff_results.values()), '-o')
|
|
plt.xticks(rotation=45)
|
|
plt.title("同一y坐标下左右液面直线x差变化趋势")
|
|
plt.ylabel("x差 (像素)")
|
|
plt.xlabel("图片序号")
|
|
plt.tight_layout()
|
|
plt.show()
|
|
print("\n=== 液面变化结果 ===")
|
|
for k, v in xdiff_results.items():
|
|
print(f"{k}: x差 = {v:.2f}")
|
|
else:
|
|
print("⚠️ 未检测到有效结果")
|
|
|
|
|
|
# ===================== 使用示例 =====================
|
|
if __name__ == "__main__":
|
|
process_image_folder(IMAGE_DIR, MODEL_PATH)
|