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

 

# -*-	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)