使用SAM分割一切模型对已有检测框的数据生成json分割标注-纯代码干货分享

Ysn0719

已于 2024-08-21 17:40:26 修改

阅读量34

点赞数 1

文章标签： json YOLO 数据库 SAM 大模型使用

于 2024-08-21 17:33:33 首次发布

本文链接：https://blog.csdn.net/yangsn0719/article/details/141399123

版权

这里默认你已经安装好了yolov8的运行环境，我这里安装的ultralytics版本是8.2.54，可以支持SAM的调用
参考：https://docs.ultralytics.com/zh/models/sam/

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# @author: ysn
# @desc：
# 功能：使用sam分割模型对已经有标注框的数据打标，生成json格式的多边形分割标注
# @date: 2024-08-20 18:24:24
# @Email：yangsn0719@163.com
# @software: vscode

import os
import cv2
import json
import math
import numpy as np
from PIL import Image
import xml.etree.ElementTree as ET
from ultralytics.models.sam import Predictor as SAMPredictor

def read_xml_obj(xml_file_path, classname):
    """读取xml文件 记录目标标签label和bbox"""
    # Parse the XML file
    tree = ET.parse(xml_file_path)
    root = tree.getroot()
    # Convert width and height to integers
    size = root.find('size')
    # width = int(float(size.find('width').text))
    # height = int(float(size.find('height').text))
    labels, bboxs = [], [] 
    # Iterate over all object elements
    for obj in root.findall('object'):
        label = obj.find('name').text
        if label in classname:
            labels.append(label)
            xmlbox = obj.find("bndbox")
            xmin, ymin, xmax, ymax = int(float(xmlbox.find("xmin").text)), int(float(xmlbox.find("ymin").text)), int(float(xmlbox.find("xmax").text)), int(float