半自动辅助制作数据集【实例分割】

news2024/11/23 12:48:59

利用yoloV8的实例分割模型，半自动辅助制作数据集

引言：【主要步骤】

步骤1：无人机航拍，收集基础图片
步骤2：将收集到的图片，全部用yoloV8-seg.pt模型进行实例分割【预测之前，将配置文件default.yaml的save_txt:设置为True】
步骤3：将txt标注文件转换为JSON格式
步骤4：打开labelme开始修改再加工标注

步骤1：【比如，如下规格】
图片特征数据：车辆、人，等等
图片拍摄变量：无人机高度、相机水平方向、相机俯仰角度
（a）固定高度、水平方向，调整俯仰角
（b）固定水平方向、俯仰角，调整高度
（c）固定俯仰角、高度，调整水平方向
步骤2：【预测代码及配置】
default.yaml【save_txt:设置为True】
使用脚本预测：my_predict.py【source参数：存放图片的目录路径】

得到的预测结果会存放在run文件夹下
步骤3：【将txt文件转换为JSON文件】
利用我在Cityscapes数据集转换为COCO数据集的文章中写的方法，改一下坐标为缩放前的即可：

import os
import cv2
import json
import glob
import numpy as np

def convert_txt_to_labelme_json(txt_path, image_path, output_dir, image_fmt='.jpg'):
    # txt 转labelme json
    txts = glob.glob(os.path.join(txt_path, "*.txt"))
    for txt in txts:
        labelme_json = {
            'version': '4.5.7',
            '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:
                content = line.split(' ')
                label = content[0]
                tem_label=str(label)
                #   0: car #车
                #   1: street  #街道、路
                #   2: person #人
                #   3: lawn #草坪
                #   4: construction #建筑
                #   5: tree #树，树林
                mapLabel={
                    "0":"car",
                    "1":"street",
                    "2":"person",
                    "3":"lawn",
                    "4":"construction",
                    "5":"tree"
                }
                shape = {
                    'label': mapLabel.get(str(label)),
                    'flags': {},
                    'points': []
                }
                for i in range(len(content)):
                    if 2 * i + 1 >= len(content):
                        break
                    else:
                        try:
                            shape['points'].append([float(content[2 * i + 1])*w, float(content[2 * i + 2])*h])
                        except Exception as e:
                            print(e)

                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=4)
            fd.close()
            print("save json={}".format(json_name_path))


if __name__=="__main__":
    in_imgs_dir = 'D:\\yoloProject\\ultralytics-registry\\test_11\\imgs'
    in_label_txt_dir = 'D:\\yoloProject\\ultralytics-registry\\test_11\\labels'
    out_labelme_json_dir = 'D:\\yoloProject\\ultralytics-registry\\test_11\\jsons'

    if not os.path.exists(out_labelme_json_dir):
        os.mkdir(out_labelme_json_dir)
    convert_txt_to_labelme_json(in_label_txt_dir,in_imgs_dir,out_labelme_json_dir,image_fmt='.png')