huangxin/zjsh_yolov11
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

bushu

Browse Source
This commit is contained in:
琉璃月光
2025-10-21 14:11:52 +08:00
parent 349449f2b7
commit df7c0730f5
363 changed files with 5386 additions and 578 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
yemian/yemian_line/tuili_c_f.py Normal file
View File

@ -0,0 +1,190 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
from pathlib import Path
from ultralytics import YOLO
# --------------------
# 参数设置（固定在脚本中）
# --------------------
INPUT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/test_l" # 图片文件夹
MODEL_PATH = "best.pt" # YOLO 模型
OUTPUT_DIR = "./output" # 保存结果
TARGET_SIZE = 640 # YOLO 输入尺寸
DIST_THRESH = 15 # 剔除离群点阈值
MAX_CORNERS = 200 # goodFeaturesToTrack 最大角点数
QUALITY_LEVEL = 0.01 # goodFeaturesToTrack qualityLevel
MIN_DISTANCE = 5 # goodFeaturesToTrack minDistance
# 全局 ROI 定义
ROIS = [
(859, 810, 696, 328), # (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 np.zeros((0,2), dtype=np.int32), np.zeros((0,2), dtype=np.int32)
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]
def filter_by_y_variation(pts):
if len(pts) < 2:
return pts.astype(np.int32)
pts_sorted = pts[np.argsort(pts[:,1])]
diffs = np.abs(np.diff(pts_sorted[:,1]))
keep_idx = np.where(diffs > y_var_thresh)[0]
if len(keep_idx) == 0:
return pts_sorted.astype(np.int32)
selected = [pts_sorted[i] for i in keep_idx] + [pts_sorted[i+1] for i in keep_idx]
selected = np.array(selected)
selected = selected[np.argsort(selected[:,1])]
_, idx = np.unique(selected.reshape(-1,2), axis=0, return_index=True)
selected = selected[np.sort(idx)]
return selected.astype(np.int32)
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=DIST_THRESH):
if pts is None or len(pts) < 2:
return None, np.zeros((0,2), dtype=np.int32)
pts = np.array(pts, dtype=np.float64)
x = pts[:,0]
y = pts[:,1]
try:
m, b = np.polyfit(y, x, 1)
except:
return None, np.zeros((0,2), dtype=np.int32)
x_fit = m*y + b
dists = np.abs(x - x_fit)
mask = dists < dist_thresh
if np.sum(mask) < 2:
return (m,b), pts.astype(np.int32)
x2, y2 = x[mask], y[mask]
try:
m2, b2 = np.polyfit(y2, x2, 1)
except:
return (m,b), np.stack([x2,y2],axis=1).astype(np.int32)
inliers = np.stack([x2,y2],axis=1).astype(np.int32)
return (m2,b2), inliers
# --------------------
# 单张图 ROI 处理
# --------------------
def process_roi_on_image(orig_img, roi):
rx, ry, rw, rh = roi
h_img, w_img = orig_img.shape[:2]
rx = max(0, rx); ry = max(0, ry)
rw = min(rw, w_img - rx); rh = min(rh, h_img - ry)
roi_img = orig_img[ry:ry+rh, rx:rx+rw].copy()
if roi_img.size == 0:
return None
resized = cv2.resize(roi_img, (TARGET_SIZE, TARGET_SIZE))
results = MODEL(resized, imgsz=TARGET_SIZE, verbose=False)
result = results[0]
overlay_roi = roi_img.copy()
if result.masks is None or len(result.masks.data)==0:
return overlay_roi
mask = result.masks.data[0].cpu().numpy()
mask_bin = (mask>0.5).astype(np.uint8)
mask_bin = cv2.resize(mask_bin,(rw,rh), interpolation=cv2.INTER_NEAREST)
# mask 半透明覆盖
color_mask = np.zeros_like(overlay_roi, dtype=np.uint8)
color_mask[mask_bin==1] = (0,255,0)
overlay_roi = cv2.addWeighted(overlay_roi,0.7,color_mask,0.3,0)
# 角点检测
mask_gray = (mask_bin*255).astype(np.uint8)
corners = cv2.goodFeaturesToTrack(mask_gray, maxCorners=MAX_CORNERS,
qualityLevel=QUALITY_LEVEL, minDistance=MIN_DISTANCE)
left_pts, right_pts = select_edge_corners(corners, rw)
left_line, left_inliers = fit_line_with_outlier_removal(left_pts)
right_line, right_inliers = fit_line_with_outlier_removal(right_pts)
# 可视化 inliers
for (cx,cy) in left_inliers:
cv2.circle(overlay_roi,(int(cx),int(cy)),4,(0,0,255),-1)
for (cx,cy) in right_inliers:
cv2.circle(overlay_roi,(int(cx),int(cy)),4,(255,0,0),-1)
# 拟合直线
if left_line is not None:
m,b = left_line
y1,y2 = 0, rh-1
x1 = int(m*y1+b); x2 = int(m*y2+b)
cv2.line(overlay_roi,(x1,y1),(x2,y2),(0,0,200),3)
if right_line is not None:
m,b = right_line
y1,y2 = 0, rh-1
x1 = int(m*y1+b); x2 = int(m*y2+b)
cv2.line(overlay_roi,(x1,y1),(x2,y2),(200,0,0),3)
return overlay_roi
# --------------------
# 批量推理文件夹
# --------------------
def infer_folder_images():
input_dir = Path(INPUT_DIR)
output_dir = Path(OUTPUT_DIR)
output_dir.mkdir(parents=True, exist_ok=True)
exts = ('*.jpg','*.jpeg','*.png','*.bmp','*.tif','*.tiff')
files = []
for e in exts:
files.extend(sorted(input_dir.glob(e)))
if len(files)==0:
print("未找到图片文件")
return
print(f"找到 {len(files)} 张图片，开始推理...")
for img_path in files:
print("-> 处理:", img_path.name)
orig_img = cv2.imread(str(img_path))
if orig_img is None:
print(" 无法读取，跳过")
continue
out_img = orig_img.copy()
for roi in ROIS:
overlay_roi = process_roi_on_image(orig_img, roi)
if overlay_roi is not None:
rx,ry,rw,rh = roi
h,w = overlay_roi.shape[:2]
out_img[ry:ry+h, rx:rx+w] = overlay_roi
save_path = output_dir / f"mask_edge_corners_{img_path.name}"
cv2.imwrite(str(save_path), out_img)
print(" 已保存 ->", save_path.name)
print("批量推理完成，结果保存在:", output_dir)
# --------------------
# 主程序
# --------------------
if __name__ == "__main__":
MODEL = YOLO(MODEL_PATH)
infer_folder_images()