最新推送

yolo11_seg/cvattoseg.py

@@ -0,0 +1,106 @@
+
+import os
+import xml.etree.ElementTree as ET
+from pathlib import Path
+
+
+def cvat_to_yolo_seg(
+    xml_path,
+    output_dir,
+    class_name_to_id=None,
+    force_class_id=None
+):
+    """
+    将 CVAT 导出的 XML（polygon / segmentation）转换为 YOLO Segmentation 格式
+
+    Args:
+        xml_path (str): CVAT 导出的 XML 文件路径
+        output_dir (str): 输出 .txt 标注文件的目录
+        class_name_to_id (dict, optional):
+            类别名到 ID 的映射
+        force_class_id (dict, optional):
+            强制指定某些类别的 ID，例如 {"yemian": 0}
+    """
+
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    # ----------------------------
+    # 自动提取类别映射
+    # ----------------------------
+    if class_name_to_id is None:
+        class_name_to_id = {}
+        labels_elem = root.find(".//labels")
+        if labels_elem is not None:
+            for idx, label in enumerate(labels_elem.findall("label")):
+                name = label.find("name").text
+                class_name_to_id[name] = idx
+        else:
+            raise RuntimeError("❌ 未找到 <labels>，请手动提供 class_name_to_id")
+
+    print(f"原始类别映射: {class_name_to_id}")
+
+    # ----------------------------
+    # 强制修改类别 ID（新增功能）
+    # ----------------------------
+    if force_class_id:
+        for name, new_id in force_class_id.items():
+            if name in class_name_to_id:
+                old_id = class_name_to_id[name]
+                class_name_to_id[name] = new_id
+                print(f"强制修改类别映射: {name} {old_id} → {new_id}")
+            else:
+                print(f"类别 {name} 不存在，跳过")
+
+    print(f"最终类别映射: {class_name_to_id}")
+
+    # ----------------------------
+    # 遍历每一张 image
+    # ----------------------------
+    for image in root.findall("image"):
+        img_name = image.get("name")
+        width = int(image.get("width"))
+        height = int(image.get("height"))
+
+        txt_path = output_dir / (Path(img_name).stem + ".txt")
+        lines = []
+
+        for polygon in image.findall("polygon"):
+            label = polygon.get("label")
+            if label not in class_name_to_id:
+                continue
+
+            class_id = class_name_to_id[label]
+            points_str = polygon.get("points")
+
+            points = []
+            for p in points_str.strip().split(";"):
+                x, y = p.split(",")
+                x = float(x) / width
+                y = float(y) / height
+                points.append(f"{x:.6f}")
+                points.append(f"{y:.6f}")
+
+            if len(points) < 6:
+                continue
+
+            lines.append(f"{class_id} " + " ".join(points))
+
+        with open(txt_path, "w", encoding="utf-8") as f:
+            f.write("\n".join(lines))
+
+    print("✅ CVAT segmentation → YOLO seg 转换完成")
+
+
+if __name__ == "__main__":
+
+    cvat_to_yolo_seg(
+        xml_path="annotations.xml",
+        output_dir="labels_seg",
+        force_class_id={
+            "yemian": 0
+        }
+    )

yolo11_seg/yolo_seg_infer_vis60.py

@@ -0,0 +1,141 @@
+import cv2
+import numpy as np
+import os
+from ultralytics import YOLO
+
+# ---------------- 配置 ----------------
+MODEL_PATH = "60seg.pt"
+IMAGE_PATH = "3.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()

yolo11_seg/yolo_seg_infer_vis61.py

@@ -0,0 +1,141 @@
+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()

yolo11_seg/yolo_seg_infer_vis—60f.py

@@ -0,0 +1,166 @@
+import cv2
+import numpy as np
+import os
+from ultralytics import YOLO
+
+# ================= 配置 =================
+MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/60seg/exp3/weights/best.pt"
+IMAGE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/分割60/class4/1"
+OUT_DIR = "./outputs"
+
+IMG_SIZE = 640
+CONF_THRES = 0.25
+
+# 多边形简化比例（点数控制核心参数）
+EPSILON_RATIO = 0.001
+
+IMG_EXTS = (".jpg", ".jpeg", ".png", ".bmp")
+
+# -------- 保存开关 --------
+SAVE_LABELS = True     # 是否保存 YOLO seg 标签
+SAVE_VIS = True       # 是否保存可视化结果
+# ======================================
+
+
+def simplify_polygon(poly, epsilon_ratio):
+    """使用 approxPolyDP 简化多边形"""
+    poly = poly.astype(np.int32)
+    perimeter = cv2.arcLength(poly, True)
+    epsilon = epsilon_ratio * perimeter
+    approx = cv2.approxPolyDP(poly, epsilon, True)
+    return approx.reshape(-1, 2)
+
+
+def extract_simplified_masks(result):
+    """
+    提取并简化 YOLO mask
+    返回: [(cls_id, poly), ...]
+    """
+    simplified = []
+
+    if result.masks is None:
+        return simplified
+
+    boxes = result.boxes
+
+    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 = simplify_polygon(poly, EPSILON_RATIO)
+
+        if len(poly) < 3:
+            continue
+
+        simplified.append((cls_id, poly))
+
+    return simplified
+
+
+def save_yolo_seg_labels(masks, img_shape, save_path):
+    """保存 YOLO segmentation 标签（无目标也生成空 txt）"""
+    h, w = img_shape[:2]
+    lines = []
+
+    for cls_id, poly in masks:
+        poly_norm = []
+        for x, y in poly:
+            poly_norm.append(f"{x / w:.6f}")
+            poly_norm.append(f"{y / h:.6f}")
+
+        lines.append(str(cls_id) + " " + " ".join(poly_norm))
+
+    with open(save_path, "w") as f:
+        if lines:
+            f.write("\n".join(lines))
+
+
+def draw_polygons(img, masks):
+    """在图像上绘制 segmentation 多边形"""
+    vis = img.copy()
+
+    for cls_id, poly in masks:
+        poly = poly.astype(np.int32)
+
+        cv2.polylines(
+            vis,
+            [poly],
+            isClosed=True,
+            color=(0, 255, 0),
+            thickness=2
+        )
+
+        x, y = poly[0]
+        cv2.putText(
+            vis,
+            str(cls_id),
+            (int(x), int(y)),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.6,
+            (0, 255, 0),
+            2
+        )
+
+    return vis
+
+
+def run_folder_inference():
+    # 输出目录
+    out_lbl_dir = os.path.join(OUT_DIR, "labels")
+    out_img_dir = os.path.join(OUT_DIR, "images")
+
+    if SAVE_LABELS:
+        os.makedirs(out_lbl_dir, exist_ok=True)
+    if SAVE_VIS:
+        os.makedirs(out_img_dir, exist_ok=True)
+
+    # 加载模型（只一次）
+    model = YOLO(MODEL_PATH)
+
+    img_files = sorted([
+        f for f in os.listdir(IMAGE_DIR)
+        if f.lower().endswith(IMG_EXTS)
+    ])
+
+    print(f"📂 共检测 {len(img_files)} 张图片")
+
+    for idx, img_name in enumerate(img_files, 1):
+        img_path = os.path.join(IMAGE_DIR, img_name)
+        img = cv2.imread(img_path)
+
+        if img is None:
+            print(f"⚠️ 跳过无法读取: {img_name}")
+            continue
+
+        results = model(
+            img,
+            imgsz=IMG_SIZE,
+            conf=CONF_THRES,
+            verbose=False
+        )
+        result = results[0]
+
+        masks = extract_simplified_masks(result)
+        base_name = os.path.splitext(img_name)[0]
+
+        # ---------- 保存标签 ----------
+        if SAVE_LABELS:
+            label_path = os.path.join(out_lbl_dir, base_name + ".txt")
+            save_yolo_seg_labels(masks, img.shape, label_path)
+
+        # ---------- 保存可视化 ----------
+        if SAVE_VIS:
+            vis_img = draw_polygons(img, masks)
+            vis_path = os.path.join(out_img_dir, img_name)
+            cv2.imwrite(vis_path, vis_img)
+
+        print(f"[{idx}/{len(img_files)}] ✅ {img_name}")
+
+    print("🎉 推理完成")
+
+
+if __name__ == "__main__":
+    run_folder_inference()

yolo11_seg/yolo_seg_infer_vis—f.py

@@ -0,0 +1,125 @@
+import cv2
+import numpy as np
+import os
+from ultralytics import YOLO
+
+# ---------------- 配置 ----------------
+MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/61seg/exp2/weights/best.pt"
+IMAGE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/cls-61/class4c/12.17-18"
+OUT_DIR = "./outputs"          # 只保存 labels
+IMG_SIZE = 640
+CONF_THRES = 0.25
+
+# 多边形简化比例（点数控制核心参数）
+EPSILON_RATIO = 0.001
+
+IMG_EXTS = (".jpg", ".jpeg", ".png", ".bmp")
+# -------------------------------------
+
+
+def simplify_polygon(poly, epsilon_ratio):
+    """
+    使用 approxPolyDP 简化多边形
+    """
+    poly = poly.astype(np.int32)
+    perimeter = cv2.arcLength(poly, True)
+    epsilon = epsilon_ratio * perimeter
+    approx = cv2.approxPolyDP(poly, epsilon, True)
+    return approx.reshape(-1, 2)
+
+
+def extract_simplified_masks(result):
+    """
+    提取并简化 YOLO mask
+    返回：
+      [(cls_id, poly), ...]
+    """
+    simplified = []
+
+    if result.masks is None:
+        return simplified
+
+    boxes = result.boxes
+
+    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 = simplify_polygon(poly, EPSILON_RATIO)
+
+        if len(poly) < 3:
+            continue
+
+        simplified.append((cls_id, poly))
+
+    return simplified
+
+
+def save_yolo_seg_labels(masks, img_shape, save_path):
+    """
+    保存 YOLO segmentation 标签
+    """
+    h, w = img_shape[:2]
+    lines = []
+
+    for cls_id, poly in masks:
+        poly_norm = []
+        for x, y in poly:
+            poly_norm.append(f"{x / w:.6f}")
+            poly_norm.append(f"{y / h:.6f}")
+
+        lines.append(str(cls_id) + " " + " ".join(poly_norm))
+
+    # ⚠️ 没有目标也生成空 txt（YOLO 训练需要）
+    with open(save_path, "w") as f:
+        if lines:
+            f.write("\n".join(lines))
+
+
+def run_folder_inference():
+    out_lbl_dir = os.path.join(OUT_DIR, "labels")
+    os.makedirs(out_lbl_dir, exist_ok=True)
+
+    # 模型只加载一次
+    model = YOLO(MODEL_PATH)
+
+    img_files = sorted([
+        f for f in os.listdir(IMAGE_DIR)
+        if f.lower().endswith(IMG_EXTS)
+    ])
+
+    print(f"📂 共检测 {len(img_files)} 张图片")
+
+    for idx, img_name in enumerate(img_files, 1):
+        img_path = os.path.join(IMAGE_DIR, img_name)
+        img = cv2.imread(img_path)
+
+        if img is None:
+            print(f"⚠️ 跳过无法读取: {img_name}")
+            continue
+
+        results = model(
+            img,
+            imgsz=IMG_SIZE,
+            conf=CONF_THRES,
+            verbose=False
+        )
+        result = results[0]
+
+        masks = extract_simplified_masks(result)
+
+        base_name = os.path.splitext(img_name)[0]
+        label_path = os.path.join(out_lbl_dir, base_name + ".txt")
+
+        save_yolo_seg_labels(masks, img.shape, label_path)
+
+        print(f"[{idx}/{len(img_files)}] ✅ {img_name}")
+
+    print("🎉 标签生成完成（仅保存 YOLO Seg 标签）")
+
+
+if __name__ == "__main__":
+    run_folder_inference()