labelme转yolo格式

时间太悠久了，不知道代码是引用的谁的了。

首先需要将labelme文件转成json文件，这个文件包含了所有的img的内容，之前的是好几个文件，

import os
import json
import numpy as np
import glob
import shutil
np.random.seed(41)
import cv2

# 0为背景
classname_to_id = {"person": 0}

class Lableme2CoCo:

    def __init__(self, splitDir=''):
        self.images = []
        self.annotations = []
        self.categories = []
        self.img_id = 0
        self.ann_id = 0
        self.splitDir = splitDir

    def save_coco_json(self, instance, save_path):
        json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)  # indent=2 更加美观显示

    # 由json文件构建COCO
    def to_coco(self, json_path_list):
        self._init_categories()
        for json_path in json_path_list:
            # print(type(json_path))
            obj = self.read_jsonfile(json_path)
            self.images.append(self._image(obj, json_path))

            shapes = obj['shapes']
            groupIds = []
            for shape in shapes:
                groupId = shape['group_id']
                groupIds.append(groupId)
            for i in set(groupIds):
                keyPoints = [0] * 51
                keyPointNum = 0
                bbox = []

                for shape in shapes:
                    if i != shape['group_id']:
                        continue
                    if shape['shape_type'] == "point":
                        labelNum = int(shape['label'])
                        keyPoints[labelNum * 3 + 0] = int(shape['points'][0][0] + 0.5)
                        keyPoints[labelNum * 3 + 1] = int(shape['points'][0][1] + 0.5)
                        keyPoints[labelNum * 3 + 2] = 2
                        keyPointNum += 1
                    if shape['shape_type'] == 'rectangle':
                        x0, y0, x1, y1 = shape['points'][0][0], shape['points'][0][1], \
                            shape['points'][1][0], shape['points'][1][1]
                        xmin = min(x0, x1)
                        ymin = min(y0, y1)
                        xmax = max(x0, x1)
                        ymax = max(y0, y1)

                        bbox = [xmin, ymin, xmax - xmin, ymax - ymin]

                annotation = self._annotation(bbox, keyPoints, keyPointNum)
                self.annotations.append(annotation)
                self.ann_id += 1
            self.img_id += 1

            # for shape in shapes:
            #     label = shape['label']
            #     if label != 'person':
            #         continue
            #
            #     annotation = self._annotation(shape)
            #     self.annotations.append(annotation)
            #     self.ann_id += 1
            # self.img_id += 1
        instance = {}
        instance['info'] = 'spytensor created'
        instance['license'] = ['license']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    # 构建类别
    def _init_categories(self):
        for k, v in classname_to_id.items():
            category = {}
            category['id'] = v
            category['name'] = k
            self.categories.append(category)

    # 构建COCO的image字段
    def _image(self, obj, jsonPath):
        image = {}
        # img_x = utils.img_b64_to_arr(obj['imageData'])
        # h, w = img_x.shape[:-1]
        jpgPath = jsonPath.replace('.json', '.jpg')
        jpgData = cv2.imread(jpgPath)
        h, w, _ = jpgData.shape

        image['height'] = h
        image['width'] = w
        image['id'] = self.img_id

        # image['file_name'] = os.path.basename(jsonPath).replace(".json", ".jpg")
        image['file_name'] = jpgPath.split(self.splitDir)[-1].replace('\\', '/')


        return image

    # 构建COCO的annotation字段
    def _annotation(self, bbox, keyPoints, keyNum):
        annotation = {}

        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        annotation['category_id'] = 1
        # annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
        annotation['segmentation'] = []
        annotation['bbox'] = bbox
        annotation['iscrowd'] = 0
        annotation['area'] = bbox[2] * bbox[3]
        annotation['keypoints'] = keyPoints
        annotation['num_keypoints'] = keyNum

        return annotation

    # 读取json文件，返回一个json对象
    def read_jsonfile(self, path):
        with open(path, "r", encoding='utf-8') as f:
            return json.load(f)

    # COCO的格式： [x1,y1,w,h] 对应COCO的bbox格式
    def _get_box(self, points):
        min_x = min_y = np.inf
        max_x = max_y = 0
        for x, y in points:
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
        return [min_x, min_y, max_x - min_x, max_y - min_y]


if __name__ == '__main__':

    labelme_path = r"D:\COCO\labelme_to_json\keypoint\train"
    print(labelme_path)
    jsonName = labelme_path.split('\\')[-1]
    saved_coco_path = r"D:\COCO\labelme_to_json\save\train"
    print(saved_coco_path)
    #####################################
    # 这个一定要注意
    # 为了方便合入coco数据， 定义截断文件的文件夹与文件名字
    splitDirFlag = 'labelMePoint\\'
    ######################################

    # 创建文件
    if not os.path.exists("%s/annotations/ " %saved_coco_path):
        os.makedirs("%s/annotations/ " %saved_coco_path)
    json_list_path = glob.glob(os.path.join(labelme_path, '*.json'))

    train_path, val_path = json_list_path, ''
    # print(train_path)
    print("train_n:", len(train_path), 'val_n:', len(val_path))

    # 把训练集转化为COCO的json格式
    l2c_train = Lableme2CoCo(splitDirFlag)
    # print(train_path)
    train_instance = l2c_train.to_coco(train_path)
    l2c_train.save_coco_json(train_instance, '%s/annotations/%s.json ' %(saved_coco_path, jsonName))

之后需要将json转成yolo格式

import json

import cv2
import pandas as pd
from PIL import Image
from collections import defaultdict

from utils import *



def convert_coco_json(cocojsonpath, savepath ,use_keypoints=False, cls91to80=True):
    """Converts COCO dataset annotations to a format suitable for training YOLOv5 models.
    Args:
        labels_dir (str, optional): Path to directory containing COCO dataset annotation files.
        use_segments (bool, optional): Whether to include segmentation masks in the output.
        use_keypoints (bool, optional): Whether to include keypoint annotations in the output.
        cls91to80 (bool, optional): Whether to map 91 COCO class IDs to the corresponding 80 COCO class IDs.
    Raises:
        FileNotFoundError: If the labels_dir path does not exist.
    Example Usage:
        convert_coco(labels_dir='../coco/annotations/', use_segments=True, use_keypoints=True, cls91to80=True)
    Output:
        Generates output files in the specified output directory.
    """
    # save_dir = make_dirs('yolo_labels')  # output directory
    save_dir = make_dirs(savepath)  # output directory
    coco80 = coco91_to_coco80_class()

    # Import json
    for json_file in sorted(Path(cocojsonpath).resolve().glob('*.json')):
        fn = Path(save_dir) / 'labels' / json_file.stem.replace('instances_', '')  # folder name
        fn.mkdir(parents=True, exist_ok=True)
        with open(json_file) as f:
            data = json.load(f)

        # Create image dict
        images = {f'{x["id"]:d}': x for x in data['images']}
        # Create image-annotations dict
        imgToAnns = defaultdict(list)
        for ann in data['annotations']:
            imgToAnns[ann['image_id']].append(ann)

        # Write labels file
        for img_id, anns in tqdm(imgToAnns.items(), desc=f'Annotations {json_file}'):
            img = images[f'{img_id:d}']
            h, w, f = img['height'], img['width'], img['file_name']

            bboxes = []
            segments = []
            keypoints = []
            for ann in anns:
                if ann['iscrowd']:
                    continue
                # The COCO box format is [top left x, top left y, width, height]
                box = np.array(ann['bbox'], dtype=np.float64)
                box[:2] += box[2:] / 2  # xy top-left corner to center
                box[[0, 2]] /= w  # normalize x
                box[[1, 3]] /= h  # normalize y
                if box[2] <= 0 or box[3] <= 0:  # if w <= 0 and h <= 0
                    continue

                cls = coco80[ann['category_id'] - 1] if cls91to80 else ann['category_id'] - 1  # class
                box = [cls] + box.tolist()
                if box not in bboxes:
                    bboxes.append(box)
                if use_keypoints and ann.get('keypoints') is not None:
                    k = (np.array(ann['keypoints']).reshape(-1, 3) / np.array([w, h, 1])).reshape(-1).tolist()
                    k = box + k
                    keypoints.append(k)

            # Write
            fname = f.split('/')[-1]
            # with open((fn / f).with_suffix('.txt'), 'a') as file:
            with open((fn / fname).with_suffix('.txt'), 'a') as file:
                for i in range(len(bboxes)):
                    if use_keypoints:
                        line = *(keypoints[i]),  # cls, box, keypoints
                    file.write(('%g ' * len(line)).rstrip() % line + '\n')
if __name__ == '__main__':
    source = 'COCO'
    cocojsonpath = r'D:\COCO\labelme_to_json\keypoint\val'
    savepath = r'D:\COCO\labelme_to_json\save\val'
    if source == 'COCO':
        convert_coco_json(cocojsonpath,  # directory with *.json
                          savepath,
                          use_keypoints=True,
                          cls91to80=True)