在做目标分类和目标检测任务,经常用到数据集和对应的标签文件。最近一段时间我做轴承表面缺陷检测任务,数据集的制作全部由手机在真实场景拍摄,由于没有考虑到手机像素,导致每张图像都有
4032
×
3024
4032\times3024
4032×3024大小,占用内存都为3M左右,这样下来,几千张图像就是40G左右,这个数据集和内存量还是不成比例的,可以说是图像太大了。所以第一步我先对图像进行不改变图像尺寸压缩内存的方式进行处理,具体的流程可以参见这篇文章:保留原图像尺寸改变像素值(压缩存储大小)。
后来经过我训练模型,发现图像的尺寸仍然太大了,因为对于检测网络而言,喂入网络一般都以
416
×
416
416\times416
416×416读入或者
640
×
640
640\times640
640×640读入,对于原图像
4032
×
3024
4032\times3024
4032×3024这个尺寸而言,训练模型存在太多的未知,比如我用YOLOX网络进行训练,训练160个epoch,最终的结果是recall趋于0,precise都为100,搜了网上的答案,一部分都是说数据集的问题。所以我打算把
4032
×
3024
4032\times3024
4032×3024的图像进行压缩。此时又出现一个问题,由于我给图像打标签,是在原图
4032
×
3024
4032\times3024
4032×3024上打的,所以如果我压缩图像,标签文件也是需要等比例压缩的。
进入正题
数据集和标签文件都是大规模的,下面的代码实现对大规模的图像和标签文件进行压缩处理,这里的压缩不是简单地等比例压缩。都知道,标签文件的信息是对应某一张图片的,所以该代码的思想就是等比例压缩读入的图像,根据 w i d t h width width和 h e i g h t height height的压缩比例调整其标签坐标框的信息。最重要的就是代码实现了::其实原理很简单,奈何我还是琢磨了很久很久~
import xml.dom.minidom
import cv2
from PIL import Image
import os
#——————————————————————————————————————————————————————
width = 0
height = 0
def convert(input_jpg, output_jpg,input_xml,output_xml):
for filename,filename1 in zip(os.listdir(input_jpg),os.listdir(input_xml)):
#for filename1 in os.listdir(input_xml):
path = input_jpg + "/" + filename#获取文件路径
print("doing picture... ", path)
ori_img = cv2.imread(path)#读取图片
height, width = ori_img.shape[:2]
# 定义缩放信息 以等比例缩放到416为例
scale = 416 / height
height = 416
width = int(width * scale)
#----------------图片等比例压缩
img = cv2.resize(ori_img, (width, height))
cv2.imwrite(output_jpg+'/'+filename, img)
#for filename in os.listdir(input_xml):
path1 = input_xml + "/" + filename1 # 获取文件路径
print("doing xml... ", path1)
dom = xml.dom.minidom.parse(path1)
root = dom.documentElement
# 读取标注目标框
objects = root.getElementsByTagName("bndbox")
for object in objects:
xmin = object.getElementsByTagName("xmin")
xmin_data = int(float(xmin[0].firstChild.data))
# xmin[0].firstChild.data =str(int(xmin1 * x))
ymin = object.getElementsByTagName("ymin")
ymin_data = int(float(ymin[0].firstChild.data))
xmax = object.getElementsByTagName("xmax")
xmax_data = int(float(xmax[0].firstChild.data))
ymax = object.getElementsByTagName("ymax")
ymax_data = int(float(ymax[0].firstChild.data))
# 更新xml
width_xml = root.getElementsByTagName("width")
width_xml[0].firstChild.data = width
height_xml = root.getElementsByTagName("height")
height_xml[0].firstChild.data = height
#----------------标签等比例压缩
xmin[0].firstChild.data = int(xmin_data * scale)
ymin[0].firstChild.data = int(ymin_data * scale)
xmax[0].firstChild.data = int(xmax_data * scale)
ymax[0].firstChild.data = int(ymax_data * scale)
# 另存更新后的文件
with open(output_xml + '/' + filename1, 'w') as f:
dom.writexml(f, addindent=' ', encoding='utf-8')
#----等比例压缩时要注释下面的代码,该代码只用于验证压缩效果
#----可以根据保存的图片验证压缩的效果
left_top = (int(xmin_data*scale), int(ymin_data*scale))
right_down= (int(xmax_data*scale), int(ymax_data*scale))
cv2.rectangle(img, left_top, right_down, (255, 0, 0), 1)
cv2.imwrite('./Annotations_new' +'/'+filename, img)
if __name__ == '__main__':
#输入路径(根据自己的路径进行设置)
input_jpg = "G:\\new\\JPG"
input_xml = "G:\\new\\XML"
#输出保存路径(根据自己的路径进行设置)
output_jpg = "G:\\new\\output_jpg"
output_xml = "G:\\new\\output_xml"
#调用函数
convert(input_jpg, output_jpg,input_xml,output_xml)
实验结果
下面是未处理的原图和标注信息。
下面就是用该种方法实现的结果图。
这是等比例压缩后的图像和标签,可以从图像看出,图片的尺寸为
554
×
416
554\times416
554×416,标签文件也发生了变化。
解决你的问题就给个赞吧~