自学内容网 自学内容网

深度学习常用脚本

# labelme2coco.py
#!/usr/bin/env python
import argparse
import collections
import datetime
import glob
import json
import os
import os.path as osp
import sys
import uuid
import cv2
import random
import imgviz
import numpy as np
import labelme
from tqdm import tqdm


# random.seed(100)  
try:
    import pycocotools.mask
except ImportError:
    print("Please install pycocotools:\n\n    pip install pycocotools\n")
    sys.exit(1)


def walk_file(path, type=".json"):
    for root, dirs, files in os.walk(path):
        for file in files:
            if file.endswith(type):
                f_path = os.path.join(root, file)
                if osp.exists(f_path) and osp.exists(f_path.replace('.json', '.jpg')):
                    yield f_path

def check_points(points):
    x0, y0 = points[0]
    x1, y1 = points[1]
    if x0 > x1:
        x0, x1 = x1, x0
    if y0 > y1:
        y0, y1 = y1, y0
    return [[x0, y0], [x1, y1]]

def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    # --input_dir  图片和labelme的路径
    # --output_dir coc的保存路径
    # labels:注意第一行是__ignore__
    # 要改
    startIndex = 1
    # label_name = ['损件', '露铜', '极反', '多件', '错料', '反向', '少锡', '侧立', '立碑', '反白', '翘脚', '短路','少件', '偏位','器件连接','引脚错位']
    label_name = ['短路', '器件连接', '少件', '偏位', '立碑','侧立', '异物', '少锡', '损件', '多件','反白', '翘脚', '引脚错位', '其他']
    # label_name = ['位移', '侧立', '反贴', '少锡', '异物', '文字错误', '极性','浮起', '短路', '立碑', '缺件', '规格错误', '锡点']
    # label_name = ['极性','丝印']
    # label_name = ['器件', '器件+焊盘']
    # label_name = ['划痕', '多墨', '墨点', '条码墨桥', '条码文字白点', '脏污']
    # label_name = ['断线','浓淡不均','墨点','划痕','破损','脏污','其它']
    # label_name = ['器件','焊盘']
    # label_name = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
    # label_name = ['反贴','正常']
    # label_name = ['anhen', 'bainingjiao', 'ganlieningjiao', 'heidian', 'liangdian', 'lianghen', 'posunningjiao', 'qipao']
    parser.add_argument("--input_dir", default=r"C:\Users\61082\Desktop\bentiiamge\jlc\1107标注任务-1730946188818", help="input annotated directory")
    parser.add_argument("--output_dir", default=r"C:\Users\61082\Desktop\bentiiamge\jlc\1107标注任务-1730946188818\annotations", help="output dataset directory")
    # parser.add_argument("--labels", default=r"D:\10-dataset\jlc\otherImage\coco_20642\1218-1702881442604\0.1\classes.txt", help="labels file")
    parser.add_argument(
        "--noviz", default=True, help="no visualizationd", action="store_true"
    )
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    print("Creating dataset:", args.output_dir)

    now = datetime.datetime.now()

    data = dict(
        licenses=[
            {
                "url":"",
                "id":0,
                "name":""
            }
        ],
        info={
            "contributor": "",
            "date_created": "",
            "description":"",
            "url":"",
            "version":"",
            "year":""
        },
        categories=[
            # supercategory, id, name
        ],

        images=[
            # license, url, file_name, height, width, date_captured, id
        ],
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
    )

    class_name_to_id = {}

    for i, line in enumerate(label_name):
        class_id = i  # starts with -1
        class_name = line
        class_name_to_id[class_name] = class_id
        data["categories"].append(
            {
                "id":class_id+startIndex,  # id从1开始
                "name":class_name,
                "supercategory":""
            }
        )

    out_ann_file = osp.join(args.output_dir, "instances_annotations.json")

    label_files = list(walk_file(args.input_dir))

    for image_id, filename in enumerate(tqdm(label_files)):
        assert os.path.exists(filename)
        # print(filename)
        label_file = json.load(open(filename, 'r',encoding='UTF-8'))   #labelme.LabelFile(filename=filename)

        base = osp.splitext(osp.basename(filename))[0]
        out_file_name = base + ".jpg"

        img_path = filename.replace(".json", ".jpg")
        assert osp.exists(img_path), img_path
        # img = cv2.imread(img_path)
        img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        data["images"].append(
            {
                "id":image_id,  # 图像id从1开始
                "width":img.shape[1],
                "height":img.shape[0],
                "file_name":out_file_name,
                "license":0,
                "flickr_url":"",
                "coco_url" :"",
                "date_captured":0,
            }
        )

        Area = []
        Bbox = []
        Labelid = []
        for shape in label_file['shapes']:
            points = shape["points"]
            label = shape["label"]

            labelid = label_name.index(label)+startIndex
            group_id = shape["group_id"]
            shape_type = shape["shape_type"]

            xmin = round(points[0][0],2)
            ymin = round(points[0][1],2)
            xmax = round(points[1][0],2)
            ymax = round(points[1][1],2)

            o_width = round(xmax-xmin,2)
            o_height = round(ymax - ymin,2)
            seg_area = o_width*o_height

            Labelid.append(labelid)
            Area.append(seg_area)
            Bbox.append([xmin,ymin,o_width,o_height])

        for area, bbox,labelid in zip(Area,Bbox,Labelid):
            data["annotations"].append(
                {
                    "id":len(data["annotations"]),
                    "image_id":image_id,
                    "category_id":labelid,
                    "segmentation":[],
                    "area":area,
                    "bbox":bbox,
                    "iscrowd":0,
                    "attributes": {"occluded": False}
                }
            )

    with open(out_ann_file, "w",encoding='UTF-8') as f:
        json.dump(data, f)

    print('finish')


if __name__ == "__main__":
    main()
    print('finish')

2、coco2labelme:

import json
import os
import os.path as osp
from pycocotools.coco import COCO
from tqdm import tqdm
import numpy as np
import cv2
#import mmcv
from PIL import Image
import shutil

# img_path = r'C:\Users\61082\Desktop\label\图像\-9223372036854773339.jpg'
# target_dir = r'C:\Users\61082\Desktop\label\imcode_img.jpg'
# shutil.copy2(img_path,target_dir)
# exit()


# 图片不存在的
no_image = []
# json不存在的
no_json = []

Image_Dir = []
CoCo_Dir = []



Image_Dir = [r'C:\Users\61082\Desktop\bentiiamge\jlc\1107标注任务-1730946188818\images']
CoCo_Dir = [r'C:\Users\61082\Desktop\bentiiamge\jlc\1107标注任务-1730946188818\annotations\instances_annotations.json']
save_dir = [r'C:\Users\61082\Desktop\bentiiamge\jlc\1107标注任务-1730946188818\labelme']

# 删除说有的labelme文件夹
for name in save_dir:
    if os.path.exists(name):
        # 删除文件夹及其中所有内容
        shutil.rmtree(name)
        print(f"文件夹 {name} 及其内容已删除")
    else:
        print(f"文件夹 {name} 不存在")


for img_dir_index in range(len(Image_Dir)):
    img_dir = Image_Dir[img_dir_index]
    os.makedirs(save_dir[img_dir_index], exist_ok=True)

    # 判断图像和json是否存在
    # 加载COCO JSON文件
    coco = COCO(CoCo_Dir[img_dir_index])

    categories = {cat['id']: cat['name'] for cat in coco.loadCats(coco.getCatIds())}
    print('categories',coco.loadCats(coco.getCatIds()))
    # exit()
    # 加载图像信息
    image_ids = coco.getImgIds()

    for img_id in tqdm(image_ids):
        img_info = coco.loadImgs(img_id)[0]
        img_path = os.path.join(img_dir, img_info['file_name'])
        if not osp.exists(img_path):
            no_image.append(img_path)
            continue
        #assert osp.exists(img_path), Image.open(img_path)
        ann_ids = coco.getAnnIds(imgIds=img_id)
        annotations = coco.loadAnns(ann_ids)
        shapes = []

        lable_name = []
        for ann in annotations:
            # 将名字转化为中文
            lable_name.append(categories[ann['category_id']])
            bbox = ann['bbox']
            shape = {
                'label': categories[ann['category_id']],
                'points': [
                    [bbox[0], bbox[1]],
                    [bbox[0] + bbox[2], bbox[1] + bbox[3]]
                ],
                'group_id': None,
                "description": "",
                'shape_type': 'rectangle',
                'flags': {}
            }
            shapes.append(shape)
        ann_json = {
                "version": "5.3.1",
                "flags": {},
                "shapes": shapes,
                "imagePath": img_info['file_name'],
                "imageData": None,
                "imageHeight": img_info['height'],
                "imageWidth": img_info['width']
        }

        image_path = os.path.join(img_dir, img_info['file_name'])
        json_save_path = os.path.join(save_dir[img_dir_index], img_info['file_name'].replace('.jpg', '.json'))
        if not os.path.exists(json_save_path):
            with open(json_save_path, 'w',encoding='utf-8') as f:
                json.dump(ann_json, f, indent=4)
            f.close()
            imagepath = json_save_path.replace('.json','.jpg')
            shutil.copy2(img_path, imagepath)
        else:
            print('new')
            new_json_save_path = json_save_path[0:-5]+'_'+str(img_id)+'.json'
            with open(new_json_save_path, 'w',encoding='utf-8') as f:
                json.dump(ann_json, f, indent=4)
            f.close()
            new_imagepath = new_json_save_path.replace('.json','.jpg')
            shutil.copy2(img_path, new_imagepath)

print("done ! ! !")
print('no_image:',no_image)

3、show_coco.py

import argparse
import os
import json
import cv2
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import shutil
from tqdm import tqdm
import sys


def show_coco(dataset_dir, output, show, category_id):
    if not os.path.exists(output):
        os.mkdir(output)
    
    annotations_path = os.path.join(dataset_dir, 'annotations', 'instances_annotations.json')
    images_dir = os.path.join(dataset_dir, 'images')
    
    annotations_data = json.load(open(annotations_path, encoding='utf-8'))

    id2class_name = {}
    for cate_item in annotations_data['categories']:
        id2class_name[cate_item['id']] = cate_item['name']

    print('id2class_name:',id2class_name)
    if category_id == -1:
        output = os.path.join(output, 'NG')
    else:
        output = os.path.join(output, id2class_name[category_id])

    if not os.path.exists(output):
        os.mkdir(output)

    id2image = {}
    image2bbox = {}
    for img_item in annotations_data['images']:
        id2image[img_item['id']] = img_item['file_name']
        image2bbox[img_item['file_name']] = []
    for ann_item in annotations_data['annotations']:
        if category_id != -1:
            if ann_item['category_id'] != category_id:
                continue
        file_name = id2image[ann_item['image_id']]
        image2bbox[file_name].append(ann_item)

    for file_name, bboxes in tqdm(image2bbox.items()):
        if len(bboxes) == 0:
            continue
        image_path = os.path.join(images_dir, file_name)
        # print(image_path)
        # image = cv2.imread(image_path)
        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)

        for ann_item in bboxes:
            x1, y1, w, h = map(int, ann_item['bbox'])
            x2 = x1 + w - 1
            y2 = y1 + h - 1
            cv2.rectangle(image, (x1, y1), (x2, y2), (0,0,255), 2)
            # cv2.putText(image, id2class_name[ann_item['category_id']], (int(x1), int(y1) - 2), 0, 2 / 3, [225, 255, 255], thickness=1, lineType=cv2.LINE_AA)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(image)
            fontStyle = ImageFont.truetype("font/simsun.ttc", 25, encoding="utf-8")
            draw.text((int(x1), int(y1) - 25), id2class_name[ann_item['category_id']], (225, 255, 255), font=fontStyle)
            image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
        
        
        if len(bboxes) > 0:
            # cv2.imwrite(os.path.join(output, file_name), image)
            cv2.imencode('.jpg', image)[1].tofile(os.path.join(output, file_name))
        if show:
            cv2.imshow('show', image)
            cv2.waitKey(0)
            cv2.destroyAllWindows()

def show_coco_ok(dataset_dir, output):
    output = os.path.join(output, 'OK')
    if not os.path.exists(output):
        os.mkdir(output)

    annotations_path = os.path.join(dataset_dir, 'annotations', 'instances_annotations.json')
    images_dir = os.path.join(dataset_dir, 'images')
    
    annotations_data = json.load(open(annotations_path, encoding='utf-8'))

    id2image = {}
    image2bbox = {}
    for img_item in annotations_data['images']:
        id2image[img_item['id']] = img_item['file_name']
        image2bbox[img_item['file_name']] = []
    for ann_item in annotations_data['annotations']:
        file_name = id2image[ann_item['image_id']]
        image2bbox[file_name].append(ann_item)

    for image_file, bbox in image2bbox.items():
        if len(bbox) > 0:
            continue
        shutil.copy(os.path.join(images_dir, image_file), output)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 目录下有2个文件夹
    parser.add_argument('--coco_dataset', default=r'D:\2-Project\9yue',type=str, help='coco dataset dir')
    parser.add_argument('--output',default=r'D:\2-Project\9yue\show', type=str, help='result output dir')
    parser.add_argument('--num_classes',default=14, type=int, help='number of classes')
    parser.add_argument('--show', action='store_true', help='show image')
    args = parser.parse_args()
    print(args.output)

    for category_id in range(1, args.num_classes+1):
        show_coco(args.coco_dataset, args.output, args.show, category_id)
    # show_coco(args.coco_dataset, args.output, args.show, -1)

    show_coco_ok(args.coco_dataset, args.output)

4、生成空的lablme:

import os
import os.path as osp
import cv2
import json
from tqdm import tqdm
import numpy as np


if __name__ == '__main__':
    img_dir = r"C:\Users\61082\Desktop\bentiiamge\jlc\1110偏移漏检\images"
    img_list = os.listdir(img_dir)
    img_list = [img for img in img_list if img.endswith('.jpg')]

    for img in tqdm(img_list):
        img_path = osp.join(img_dir, img)
        image = cv2.imdecode(np.fromfile(img_path), 1)
        # image = cv2.imread(img_path)
        h, w = image.shape[:2]
        # print(image.shape)

        shapes = []
        ann_json = {
            "version": "5.3.1",
            "flags": {},
            "shapes": shapes,
            "imagePath": img,
            "imageData": None,
            "imageHeight": h,
            "imageWidth": w
        }

        json_save_path = osp.join(img_dir, img.replace('.jpg', '.json'))
        if osp.exists(json_save_path):
            continue
        with open(json_save_path, 'w') as f:
            json.dump(ann_json, f, indent=4)
        f.close()

5、labelme2xyxy

import os
import json
import cv2
import base64
import numpy as np

def cv2_base64(image):
    base64_str = cv2.imencode('.jpg',image)[1].tobytes()
    base64_str = base64.b64encode(base64_str).decode()
    return base64_str

# 图片的路径
src_dir = r'D:\10-dataset\jlc\otherImage\coco_20642\1218-1702881442604\0.1\labelme'
# 生成图片和lxyxy的路径
out_dir = r'D:\10-dataset\jlc\otherImage\coco_20642\1218-1702881442604\0.1\lxylxy'
# classes.txt 的路径
cates_path = r'D:\10-dataset\jlc\otherImage\coco_20642\1218-1702881442604\0.1\classes.txt'

assert os.path.exists(src_dir)
assert os.path.exists(cates_path), '没有发现类别文件classes.txt'
os.makedirs(out_dir, exist_ok=True)

# with open(cates_path, 'r', encoding='utf-8') as f:
#     cates = f.read().strip().split('\n')

# print(cates)
#cates2id = {k:i for i,k in enumerate(cates)}
cates2id = {'侧立':1,'偏位':2,'反向':3,'反白':4,'多件':5,'少件':6,'少锡':7,'损件':8,'极反':9,'短路':10,'立碑':11,'翘脚':12,'错料':13,'露铜':14}

image_name_list = [i for i in os.listdir(src_dir) if i.endswith('.jpg')]

for name in image_name_list:
    name_path = os.path.join(src_dir, name)
    anno_path = os.path.join(src_dir, name[:-3]+'json')
    print(name_path)
    image = cv2.imread(name_path, cv2.IMREAD_COLOR)
    h, w, c = image.shape

    if not os.path.exists(anno_path):
        # cv2.imwrite(os.path.join(label_img_dir, name), image)
        continue

    with open(anno_path, 'r', encoding='utf8') as f:
        content = json.load(f)
    
   
    annos = content['shapes']
    shapes = []

    boxes = []
    for anno in annos:
        if anno['label'] not in cates2id:
            continue

        labelid = int(cates2id[anno['label']])
        labelname = anno['label']

        points = np.array(anno['points']).round().astype(int)
        box = cv2.boundingRect(points)

        boxes.append([labelid, box[0], box[1], box[0] + box[2], box[1] + box[3]])
    
    if len(boxes):
        cv2.imwrite(os.path.join(out_dir, name), image)
        with open(os.path.join(out_dir, name[:-3]+'txt'), 'w') as f:
            f.write('\n'.join([' '.join(line) for line in np.array(boxes, dtype=str).tolist()]))
    else:
        print(name_path, '\t空标注')            

6、lxyxy2labelme.py

import os
import json
import cv2
import base64

def cv2_base64(image):
    base64_str = cv2.imencode('.jpg',image)[1].tobytes()
    base64_str = base64.b64encode(base64_str).decode()
    return base64_str

# def base64_cv2(base64_str):
#     imgString = base64.b64decode(base64_str)
#     nparr = np.fromstring(imgString,np.uint8)  
#     image = cv2.imdecode(nparr,cv2.IMREAD_COLOR)
#     return image

# 图片路径
src_dir = r'D:\2-Project\4-jixing\dataset\test\images' # input('请输入目录: ')
# 生成的labelme文件路径
out_dir = src_dir + '_out' # os.path.join(base_dir, 'labelimg')
# class.txt路径
cates_path = r'D:\2-Project\4-jixing\dataset\test\classes.txt' #os.path.join(base_dir, 'classes.txt')

assert os.path.exists(src_dir)
assert os.path.exists(cates_path), '没有发现类别文件classes.txt'
os.makedirs(out_dir, exist_ok=True)

with open(cates_path, 'r',encoding='utf8') as f:
    cates = f.read().strip().split('\n')

print(cates)

image_name_list = [i for i in os.listdir(src_dir) if i.endswith('.jpg')]

for name in image_name_list:
    name_path = os.path.join(src_dir, name)
    anno_path = os.path.join(src_dir, name[:-3]+'txt')
    print(name_path)
    image = cv2.imread(name_path, cv2.IMREAD_COLOR)
    h, w, c = image.shape

    if not os.path.exists(anno_path):
        # cv2.imwrite(os.path.join(label_me_dir, name), image)
        continue

    with open(anno_path, 'r') as f:
        annos = [line.strip().split(' ') for line in f.readlines()]
    
    shapes = []

    for anno in annos:
        labelid = int(anno[0])
        labelname = cates[labelid]

        xmin = float(anno[1])
        ymin = float(anno[2])
        xmax = float(anno[3])
        ymax = float(anno[4])

        points = [
            [xmin, ymin],
            [xmax, ymax]
        ]

        shapes.append(
            {
                'label':labelname,
                'points':points,
                'group_id':None,
                'description':'',
                "shape_type": "rectangle",
                "flags": {}
            }
        )
    
    root = {
        "version": "5.2.1",
        "flags": {},
        'shapes': shapes,
        'imagePath':name,
        "imageHeight": h,
        "imageWidth": w,
        "imageData": cv2_base64(image)
    }

    if len(shapes):
        cv2.imwrite(os.path.join(out_dir, name), image)
        with open(os.path.join(out_dir, name[:-3]+'json'), 'w') as f:
            f.write(json.dumps(root, indent=4, separators=(',', ': ')))
    else:
        print(name_path, '\t空标注')

7、labelme2voc.py

import json
from xml.etree.ElementTree import Element, SubElement, tostring
from xml.dom import minidom
from lxml import etree, objectify
import argparse
import os
from tqdm import tqdm

def create_version_tree(version=''):
    version_tree = etree.Element("version")
    version_tree.text = version
    return version_tree

def labelme_to_voc(labelme_dir, voc_dir):
    files = os.listdir(labelme_dir)
    for filename in tqdm(files):
        if filename.endswith('.json'):
            labelme_path = os.path.join(labelme_dir,filename)
            with open(labelme_path, 'r',encoding="utf-8") as f:
                data = json.load(f)
            # 创建VOC格式的XML根元素
            E = objectify.ElementMaker(annotate=False)
            annotations_tree = E.annotations()
            version_tree = create_version_tree()
            annotations_tree.append(version_tree)

            folder = etree.Element('folder')
            folder.text = ""
            annotations_tree.append(folder)

            file_name = etree.Element('filename')
            file_name.text = data['imagePath']
            annotations_tree.append(file_name)

            path = etree.Element('path')
            path.text = ''
            annotations_tree.append(path)

            source = etree.Element('source')
            database = etree.Element('database')
            database.text = 'Unknown'
            source.append(database)
            annotations_tree.append(source)

            size = etree.Element('size')
            width = etree.Element('width')
            width.text = str(data['imageWidth'])
            height = etree.Element('height')
            height.text = str(data['imageHeight'])
            depth = etree.Element('depth')
            depth.text = '3'
            size.append(width)
            size.append(height)
            size.append(depth)
            annotations_tree.append(size)

            segmented = etree.Element('segmented')
            segmented.text = '0'
            annotations_tree.append(segmented)
            # 处理对象标注
            for shape in data['shapes']:
                label = shape['label']
                points = shape['points']
                if len(points) > 0:
                    object_ = etree.Element('object')
                    name = etree.Element('name')
                    name.text = label
                    pose = etree.Element('pose')
                    pose.text = 'Unspecified'
                    truncated = etree.Element('truncated')
                    truncated.text = '0'
                    difficult = etree.Element('difficult')
                    difficult.text = '0'
                    bndbox = etree.Element('bndbox')
                    xmin = etree.Element('xmin')
                    ymin = etree.Element('ymin')
                    xmax = etree.Element('xmax')
                    ymax = etree.Element('ymax')
                    xmin.text = str(int(points[0][0]))
                    ymin.text = str(int(points[0][1]))
                    xmax.text = str(int(points[1][0]))
                    ymax.text = str(int(points[1][1]))
                    bndbox.append(xmin)
                    bndbox.append(ymin)
                    bndbox.append(xmax)
                    bndbox.append(ymax)
                    object_.append(name)
                    object_.append(pose)
                    object_.append(truncated)
                    object_.append(difficult)
                    object_.append(bndbox)
                    annotations_tree.append(object_)

            # 创建一个树结构
            tree = etree.ElementTree(annotations_tree)
            # 将树结构写入到XML文件中
            voc_path = os.path.join(voc_dir, filename.replace('.json','.xml'))
            etree.ElementTree(annotations_tree).write(voc_path, encoding="utf-8", pretty_print=True)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--labelme_dir", default=r"D:\export_imgs_0205\labelme", help="input annotated directory")
    parser.add_argument("--voc_dir", default=r"D:\export_imgs_0205\voc", help="output dataset directory")
    args = parser.parse_args()
    os.makedirs(args.voc_dir, exist_ok=True)
    labelme_to_voc(args.labelme_dir, args.voc_dir)
    print('finish')

8、voc2labelme.py

import os
import xml.etree.ElementTree as ET
import json
import argparse
from tqdm import tqdm

def voc_to_labelme(voc_dir, labelme_dir):
    files = os.listdir(voc_dir)
    for filename in tqdm(files):
        if filename.endswith('.xml'):
            xml_path = os.path.join(voc_dir, filename)
            tree = ET.parse(xml_path)
            root = tree.getroot()
            image_name = filename.replace('.xml','.jpg')
            labelme_dict = {
                "version":"5.3.1",
                "flags":{},
                "shapes":[],
                "imagePath": image_name,
                "imageData": None,
                "imageHeight": 0,
                "imageWidth": 0,
            }

            # 读取图像尺寸信息
            size = root.find('size')
            labelme_dict['imageHeight'] = int(size.find('height').text)
            labelme_dict['imageWidth'] = int(size.find('width').text)

            # 读取物体标注信息
            for obj in root.findall('object'):
                labelme_obj = {
                                "label": "",
                                "points": [[0,0],[0,0]],
                                "group_id": None,
                                "description": "",
                                "shape_type": "rectangle",
                                "flags":{}
                                }

                labelme_obj['label'] = obj.find('name').text

                bndbox = obj.find('bndbox')
                labelme_obj['points'][0][0]= int(bndbox.find('xmin').text)
                labelme_obj['points'][0][1] = int(bndbox.find('ymin').text)
                labelme_obj['points'][1][0] = int(bndbox.find('xmax').text)
                labelme_obj['points'][1][1] = int(bndbox.find('ymax').text)

                labelme_dict['shapes'].append(labelme_obj)

            json_file_path = os.path.join(labelme_dir, filename[:-4] + '.json')
            with open(json_file_path, 'w',encoding='utf-8') as f:
                json.dump(labelme_dict, f, indent=4)
            f.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--voc_dir", default=r"D:\export_imgs_0205\voc", help="input annotated directory")
    parser.add_argument("--label_dir", default=r"D:\export_imgs_0205\labelme", help="output dataset directory")
    args = parser.parse_args()
    os.makedirs(args.voc_dir, exist_ok=True)
    voc_to_labelme(args.voc_dir,args.label_dir)
    print('finish')


原文地址:https://blog.csdn.net/weixin_43940118/article/details/143691664

免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!