96 lines
2.2 KiB
Python
96 lines
2.2 KiB
Python
import cv2
|
||
import numpy as np
|
||
import os
|
||
from ultralytics import YOLO
|
||
|
||
# ---------------- 配置 ----------------
|
||
MODEL_PATH = "60seg.pt"
|
||
IMAGE_PATH = "1.png" # 支持任意路径
|
||
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):
|
||
"""
|
||
仅填充 mask，不画标签、不画轮廓
|
||
同时返回每个 mask 的类别名和面积
|
||
"""
|
||
overlay = frame.copy()
|
||
|
||
if result.masks is None:
|
||
return frame, []
|
||
|
||
boxes = result.boxes
|
||
areas = []
|
||
|
||
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)
|
||
|
||
# 计算面积（像素）
|
||
area = cv2.contourArea(poly)
|
||
areas.append((result.names[cls_id], area))
|
||
|
||
# 仅填充 mask
|
||
cv2.fillPoly(overlay, [poly], color)
|
||
|
||
blended = cv2.addWeighted(overlay, ALPHA, frame, 1 - ALPHA, 0)
|
||
return blended, areas
|
||
|
||
|
||
def run_image_inference():
|
||
# 加载模型
|
||
model = YOLO(MODEL_PATH)
|
||
|
||
# 读取图片
|
||
img = cv2.imread(IMAGE_PATH)
|
||
if img is None:
|
||
raise FileNotFoundError(f" 无法读取图片: {IMAGE_PATH}")
|
||
|
||
print(f"正在处理: {IMAGE_PATH}")
|
||
|
||
# 推理
|
||
results = model(
|
||
img,
|
||
imgsz=IMG_SIZE,
|
||
conf=CONF_THRES,
|
||
verbose=False
|
||
)
|
||
result = results[0]
|
||
|
||
# 生成可视化图（无标签）
|
||
vis, areas = draw_segmentation(img, result)
|
||
|
||
# 输出到原图所在文件夹
|
||
base_name = os.path.splitext(IMAGE_PATH)[0] # 如 "/path/to/1"
|
||
out_path = base_name + "_seg.png" # → "/path/to/1_seg.png"
|
||
|
||
# 保存
|
||
cv2.imwrite(out_path, vis)
|
||
|
||
# 打印面积
|
||
print("\n📊 Mask 面积统计 (像素):")
|
||
if areas:
|
||
for name, area in areas:
|
||
print(f" • {name}: {area:.1f}")
|
||
else:
|
||
print(" (无有效 mask)")
|
||
|
||
print(f"\n完成！可视化结果已保存至:\n {out_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_image_inference() |