142 lines
3.2 KiB
Python
142 lines
3.2 KiB
Python
import cv2
|
||
import numpy as np
|
||
import os
|
||
from ultralytics import YOLO
|
||
|
||
# ---------------- 配置 ----------------
|
||
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/61seg/exp/weights/best.pt"
|
||
IMAGE_PATH = "2.png"
|
||
OUT_DIR = "outputs"
|
||
IMG_SIZE = 640
|
||
CONF_THRES = 0.25
|
||
ALPHA = 0.5
|
||
# -------------------------------------
|
||
|
||
|
||
def get_color(idx):
|
||
np.random.seed(idx)
|
||
return tuple(int(x) for x in np.random.randint(0, 255, 3))
|
||
|
||
|
||
def draw_segmentation(frame, result):
|
||
"""
|
||
原样保留你的可视化逻辑
|
||
"""
|
||
overlay = frame.copy()
|
||
|
||
if result.masks is None:
|
||
return frame
|
||
|
||
boxes = result.boxes
|
||
names = result.names
|
||
|
||
for i, poly in enumerate(result.masks.xy):
|
||
cls_id = int(boxes.cls[i])
|
||
conf = float(boxes.conf[i])
|
||
|
||
if conf < CONF_THRES:
|
||
continue
|
||
|
||
color = get_color(cls_id)
|
||
poly = poly.astype(np.int32)
|
||
|
||
# 填充 mask
|
||
cv2.fillPoly(overlay, [poly], color)
|
||
|
||
# 轮廓
|
||
cv2.polylines(overlay, [poly], True, color, 2)
|
||
|
||
# 标签
|
||
x, y = poly[0]
|
||
label = f"{names[cls_id]} {conf:.2f}"
|
||
cv2.putText(
|
||
overlay,
|
||
label,
|
||
(x, max(y - 5, 20)),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.6,
|
||
color,
|
||
2
|
||
)
|
||
|
||
return cv2.addWeighted(overlay, ALPHA, frame, 1 - ALPHA, 0)
|
||
|
||
|
||
def save_masks_as_yolo_seg(result, img_shape, save_path):
|
||
"""
|
||
保存 YOLO segmentation 标注（txt）
|
||
格式：
|
||
class_id x1 y1 x2 y2 ...（全部归一化）
|
||
"""
|
||
if result.masks is None:
|
||
return
|
||
|
||
h, w = img_shape[:2]
|
||
boxes = result.boxes
|
||
|
||
lines = []
|
||
|
||
for i, poly in enumerate(result.masks.xy):
|
||
cls_id = int(boxes.cls[i])
|
||
conf = float(boxes.conf[i])
|
||
|
||
if conf < CONF_THRES:
|
||
continue
|
||
|
||
poly_norm = []
|
||
for x, y in poly:
|
||
poly_norm.append(f"{x / w:.6f}")
|
||
poly_norm.append(f"{y / h:.6f}")
|
||
|
||
line = str(cls_id) + " " + " ".join(poly_norm)
|
||
lines.append(line)
|
||
|
||
if lines:
|
||
with open(save_path, "w") as f:
|
||
f.write("\n".join(lines))
|
||
|
||
|
||
def run_image_inference():
|
||
os.makedirs(OUT_DIR, exist_ok=True)
|
||
|
||
# 加载模型
|
||
model = YOLO(MODEL_PATH)
|
||
|
||
# 读取图片
|
||
img = cv2.imread(IMAGE_PATH)
|
||
if img is None:
|
||
raise FileNotFoundError(f"❌ 无法读取图片: {IMAGE_PATH}")
|
||
|
||
# 推理
|
||
results = model(
|
||
img,
|
||
imgsz=IMG_SIZE,
|
||
conf=CONF_THRES,
|
||
verbose=False
|
||
)
|
||
|
||
result = results[0]
|
||
|
||
# 1️⃣ 保存可视化结果
|
||
vis = draw_segmentation(img, result)
|
||
out_img_path = os.path.join(
|
||
OUT_DIR,
|
||
os.path.splitext(os.path.basename(IMAGE_PATH))[0] + "_seg.png"
|
||
)
|
||
cv2.imwrite(out_img_path, vis)
|
||
|
||
# 2️⃣ 保存 YOLO segmentation 标注
|
||
out_txt_path = os.path.join(
|
||
OUT_DIR,
|
||
os.path.splitext(os.path.basename(IMAGE_PATH))[0] + ".txt"
|
||
)
|
||
save_masks_as_yolo_seg(result, img.shape, out_txt_path)
|
||
|
||
print("✅ 推理完成")
|
||
print(f"🖼 可视化结果: {out_img_path}")
|
||
print(f"📄 YOLO Seg 标注: {out_txt_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_image_inference()
|