voc的xml,按照各个标签保存成单独的图片

#!/usr/bin/env python
#-*- coding: utf-8 -*-
# 把voc xmlæ ¼å¼è½¬æ¢ä¸ºtxtæ ¼å¼. 自己部门标注的人脸数æ®ï¼Œéœ€è¦è½¬æ?
# 目录结构如下
# subDir
#  --annotations
#  --Images

import os
import sys
import time
import random
import argparse
import shutil
import argparse
from xml.dom import minidom 
import cv2
from multiprocessing.dummy import Pool as ThreadPool
try: 
    import xml.etree.cElementTree as ET 
except ImportError: 
    import xml.etree.ElementTree as ET 

classes         = ['label1','label2','label3']
count_class_num = {} 
def GetImgNameByEveryDir(file_dir,videoProperty):  
    FileNameWithPath ,FileName ,FileDir          = [],[],[]
    for root, dirs, files in os.walk(file_dir):  
        for file in files:  
            if os.path.splitext(file)[1] in videoProperty:  
                FileNameWithPath.append(os.path.join(root, file))  # ä¿å­˜å›¾ç‰‡è·¯å¾„
                FileName.append(file)                              # ä¿å­˜å›¾ç‰‡å称
                FileDir.append(root[len(file_dir):])               # ä¿å­˜å›¾ç‰‡æ‰€åœ¨æ–‡ä»¶å¤¹
    return FileName,FileNameWithPath,FileDir

def GetBoxInfo(xmlfile,img,filename,ProcessDir):
    try:
        tree =ET.parse(xmlfile) 
        root = tree.getroot()  
        size = root.find('size')  
        w = int(size.find('width').text)  
        h = int(size.find('height').text)
    except:
        return False,0
    else:
        tree =ET.parse(xmlfile) 
        root = tree.getroot()  
        size = root.find('size')  
        w = int(size.find('width').text)  
        h = int(size.find('height').text)


    img_num = 1
    for obj in root.findall('object'):
        print('label = ',obj.find('name').text)
        if obj.find('name').text in count_class_num:
            count_class_num[obj.find('name').text]+=1
        else:
            count_class_num[obj.find('name').text]=1


        if obj.find('name').text not in classes:
            continue
        save_dir = ProcessDir + '-' + obj.find('name').text + '/'

        if os.path.exists(save_dir)==False:
            os.makedirs(save_dir)

        xmlbox = obj.find('bndbox')  
        xmin = int(float(xmlbox.find('xmin').text))     
        ymin = int(float(xmlbox.find('ymin').text))    
        xmax = int(float(xmlbox.find('xmax').text))    
        ymax = int(float(xmlbox.find('ymax').text))

        try:
            rect_img = img[ymin:ymax,xmin:xmax,:]
            cv2.imwrite(save_dir + filename + '_' + str(img_num) + '.jpg',rect_img)
        except:
            continue
        img_num+=1

        # raise


    if len(root)==0:
        return 0

    return 1

def Process(ProcessDir):
    xmlDirs     = ProcessDir #+ 'xml/' 
    imgDirs     = ProcessDir #+ 'jpg/'
    FileName1,FileNameWithPath1,FileDir1 = GetImgNameByEveryDir(xmlDirs,['.jpg','.jpeg','.png','.JPG','.JPEG','.PNG'])


    for k in range(len(FileName1)):
        print('FileName1 = ',FileName1[k])
        imgfile = FileNameWithPath1[k]
        xmlfile = os.path.split(FileNameWithPath1[k])[0] + '/' + os.path.splitext(FileName1[k])[0] + '.xml'
        img     = cv2.imread(imgfile)
        result  = GetBoxInfo(xmlfile,img,FileName1[k][:-4],ProcessDir)

            
if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--ProcessDir',type=str, default='test_dir/') 
    args = parser.parse_args()
    ProcessDir      = args.ProcessDir
    Process(ProcessDir)
    print('count_class_num = ',count_class_num)

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值