#!/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)
voc的xml,按照各个标签保存成单独的图片
最新推荐文章于 2024-10-05 21:16:56 发布