Yolo分割标签类型转Sam标签类型

Yolo2Labelme.py

import os
import glob
import numpy as np
import cv2
import json
def convert_txt_to_labelme_json(txt_path, image_path, output_dir, class_name, image_fmt='.jpg'):
    txts = glob.glob(os.path.join(txt_path, "*.txt"))
    for txt in txts:
        labelme_json = {
            'version': '5.5.0',
            'flags': {},
            'shapes': [],
            'imagePath': None,
            'imageData': None,
            'imageHeight': None,
            'imageWidth': None,
        }
        txt_name = os.path.basename(txt)
        image_name = txt_name.split(".")[0] + image_fmt
        labelme_json['imagePath'] = image_name
        image_name = os.path.join(image_path, image_name)
        if not os.path.exists(image_name):
            raise Exception('txt 文件={},找不到对应的图像={}'.format(txt, image_name))
        image = cv2.imdecode(np.fromfile(image_name, dtype=np.uint8), cv2.IMREAD_COLOR)
        h, w = image.shape[:2]
        labelme_json['imageHeight'] = h
        labelme_json['imageWidth'] = w
        with open(txt, 'r') as t:
            lines = t.readlines()
            for line in lines:
                point_list = []
                content = line.split(' ')
                label = class_name[int(content[0])]
                for index in range(1, len(content)):
                    if index % 2 == 1:
                        x = (float(content[index])) * w
                        point_list.append(x)
                    else:
                        y = (float(content[index])) * h
                        point_list.append(y)
                point_list = [point_list[i:i + 2] for i in range(0, len(point_list), 2)]
                shape = {
                    'label': label,
                    'points': point_list,
                    'group_id': None,
                    'description': None,
                    'shape_type': 'polygon',
                    'flags': {},
                    'mask': None
                }
                labelme_json['shapes'].append(shape)
            json_name = txt_name.split('.')[0] + '.json'
            json_name_path = os.path.join(output_dir, json_name)
            fd = open(json_name_path, 'w')
            json.dump(labelme_json, fd, indent=2)
            fd.close()
            print("save json={}".format(json_name_path))
if __name__ == '__main__':
    txt_path = 'labels'
    image_path = 'images'
    output_dir = 'jsons'
    class_name = ['dog' ,'cat']
    convert_txt_to_labelme_json(txt_path, image_path, output_dir, class_name)

Labelme2Sam.py

import json
import numpy as np
from skimage.draw import polygon2mask
from pycocotools.mask import encode
import os
from tqdm import tqdm
def labelme_to_coco_sam(labelme_json_path, output_json_path, category_mapping=None):
    with open(labelme_json_path, 'r') as f:
        labelme_data = json.load(f)
    coco_sam_data = {
        "image": {
            "file_name": os.path.basename(labelme_data["imagePath"]),
            "height": labelme_data["imageHeight"],
            "width": labelme_data["imageWidth"],
            "id": 0
        },
        "annotations": []
    }
    if category_mapping is None:
        categories = list({shape["label"] for shape in labelme_data["shapes"]})
        category_mapping = {name: i + 1 for i, name in enumerate(categories)}
    for i, shape in enumerate(labelme_data["shapes"]):
        if shape["shape_type"] != "polygon":
            continue
        points = np.array(shape["points"])
        height, width = labelme_data["imageHeight"], labelme_data["imageWidth"]
        mask = polygon2mask((height, width), points).astype(np.uint8)
        rle = encode(np.asfortranarray(mask))
        rle["counts"] = rle["counts"].decode("utf-8")  # bytes -> str
        x_min, y_min = np.min(points, axis=0)
        x_max, y_max = np.max(points, axis=0)
        bbox = [x_min, y_min, x_max - x_min, y_max - y_min]
        coco_sam_data["annotations"].append({
            "id": i,
            "category_id": category_mapping[shape["label"]],
            "segmentation": rle,
            "bbox": bbox,
            "area": float(np.sum(mask)),
            "iscrowd": 0
        })
    with open(output_json_path, 'w') as f:
        json.dump(coco_sam_data, f, indent=2)
def batch_convert_labelme_to_coco_sam(input_dir, output_dir, category_mapping=None):
    os.makedirs(output_dir, exist_ok=True)
    json_files = [f for f in os.listdir(input_dir) if f.endswith('.json')]
    for json_file in tqdm(json_files, desc="Converting files"):
        input_path = os.path.join(input_dir, json_file)
        output_path = os.path.join(output_dir, json_file)
        labelme_to_coco_sam(input_path, output_path, category_mapping)
if __name__ == '__main__':
    input_dir = "labelmejsons"
    output_dir = "sam1jsons"
    category_mapping = {"dog": 1, "cat": 2}
    batch_convert_labelme_to_coco_sam(input_dir, output_dir, category_mapping)
    print(f"转换完成！结果保存在 {output_dir}")