欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

【工具脚本】将VOC格式XML标注文件的目标裁剪为小图保存下来

程序员文章站 2022-04-17 21:21:33
...

 

# -*-	coding:utf-8	-*-
from xml.dom.minidom import parse
import matplotlib.pyplot as plt
import xml.dom.minidom
import os, shutil
import matplotlib
import numpy as np
import cv2
from PIL import Image, ImageDraw

##########################################################
root = "/data_1/"
# only need to change these
##########################################################
annroot = root + 'Annotations/'
picroot = root + 'JPEGImages/'
anns = os.listdir(annroot)
imgs = os.listdir(picroot)

labelmap = ["chepai", "chedeng", "chebiao","person"]

colormap = ["red", "green", "yellow", "blue"]


def mkdir(path):
    folder = os.path.exists(path)
    if not folder:
        os.makedirs(path)


# draw.rectangle((xmin, ymin, xmax, ymax), outline = colormap[i])

number = 0
nn = 0

for ann in anns:
    number += 1
    print (number)
    print (ann)
    annpath = annroot + ann
    picpath = picroot + ann.replace("xml", "jpg")
    im = Image.open(picpath)
    img = cv2.imread(picpath)
    draw = ImageDraw.Draw(im)
    DOMTree = xml.dom.minidom.parse(annpath)
    collection = DOMTree.documentElement
    objects = collection.getElementsByTagName("object")
    labelsss = ""
    for object_ in objects:
        # print (object_)
        a = object_.getElementsByTagName("name")[0].childNodes[0].nodeValue
        k = a.split('.', 1)
        kk = k[0]
        b = str(kk)
        for i in range(0, len(labelmap)):
            label = labelmap[i]
            # print (label)

            if b == label:
                nn += 1
                print (label)
                if label not in labelsss:
                    labelsss += label + "_"
                bndboxs = object_.getElementsByTagName("bndbox")
                for bndbox in bndboxs:
                    xmin = bndbox.getElementsByTagName('xmin')[0].childNodes[0].nodeValue
                    ymin = bndbox.getElementsByTagName('ymin')[0].childNodes[0].nodeValue
                    xmax = bndbox.getElementsByTagName('xmax')[0].childNodes[0].nodeValue
                    ymax = bndbox.getElementsByTagName('ymax')[0].childNodes[0].nodeValue
                    xtmp1 = xmin.split('.', 1)
                    xmin1 = xtmp1[0]
                    xtmp2 = xmax.split('.', 1)
                    xmax1 = xtmp2[0]
                    xtmp3 = ymin.split('.', 1)
                    ymin1 = xtmp3[0]
                    xtmp4 = ymax.split('.', 1)
                    ymax1 = xtmp4[0]
                    xmin = int(xmin1)
                    ymin = int(ymin1)
                    xmax = int(xmax1)
                    ymax = int(ymax1)
                    if xmin < 0:
                        xmin = 0
                    if ymin < 0:
                        ymin = 0
                    sp = img.shape
                    if xmax > sp[1]:
                        xmax = sp[1]
                    if ymax > sp[0]:
                        ymax = sp[0]
                    draw.rectangle((xmin, ymin, xmax, ymax), outline=colormap[i])
                    draw.rectangle((xmin - 1, ymin - 1, xmax - 1, ymax - 1), outline=colormap[i])
                    draw.rectangle((xmin + 1, ymin + 1, xmax + 1, ymax + 1), outline=colormap[i])
                #    draw.rectangle((xmin-2, ymin-2, xmax-2, ymax-2), outline = colormap[i])
                #    draw.rectangle((xmin+2, ymin+2, xmax+2, ymax+2), outline = colormap[i])
                #    draw.rectangle((xmin-3, ymin-3, xmax-3, ymax-3), outline = colormap[i])
                #    draw.rectangle((xmin+3, ymin+3, xmax+3, ymax+3), outline = colormap[i])

                roiimg = img[ymin: ymax, xmin:xmax]
                save_op = root + '0_xiaotu/' + label + "_" + "/"
                mkdir(save_op)
                saveopath = save_op + str(nn) + "_" + ann.replace("xml", "jpg")
                if (roiimg.shape[0] <= 0 or roiimg.shape[1] <= 0):
                    with open("size_error.txt", "a") as f:
                        f.write("\n{}\n the img width or hight is error!\n".format(ann))  # 这句话自带文件关闭功能,不需要再写f.close()
                else:
                    cv2.imwrite(saveopath, roiimg)

                # break
        label_has = 0
        for label in labelmap:
            if b != label:
                label_has = 1
        if not label_has:
            print (ann + "======" + b + "============================")

    save_p = root + '0_check/' + labelsss + "/"
    savepath = save_p + ann.replace("xml", "jpg")
    mkdir(save_p)
    im.save(savepath)
    # cv2.imwrite(savepath,roiimg)

 

相关标签: 工具脚本