遇到的项目,需要对数据集中的图片和框转化大小,因此为了看转化的数据集是否正确,写下如下代码分享给大家:
import numpy as np
import xml.etree.ElementTree as ET
import os
import torch.cuda
import cv2
import glob
from PIL import Image
import time
def pali_box(img_trnsor,bbox_tensor):
'''
Args:
img_trnsor:输入图片为tensor(1,C,H,W)
bbox_tensor: 框的位置(ymin,xmin,ymax,xmax)
Returns:
注意.copy()的使用
显示带有框的图片
'''
img = torch.squeeze(img_trnsor)
img = torch.permute(img, (1, 2, 0)).numpy()
bbox_1 = bbox_tensor
box_color = (0, 0, 255)
img1 = np.array(img, dtype='uint8').copy()
img1 = img1[:,:,::-1].copy()
for bbox_ in bbox_1:
ymin, xmin, ymax, xmax = int(bbox_[0]), int(bbox_[1]), int(bbox_[2]), int(bbox_[3])
cv2.rectangle(img1, (xmin, ymin), (xmax, ymax), color=box_color, thickness=4)
cv2.imshow('image', img1)
cv2.waitKey(1000) == ord('q') # 按'q'键推出
cv2.destroyAllWindows()
def box(path):
'''
Args:
path:xml文件地址
Returns:
返回xml中的图片四个坐标(ymin,xmin,ymax,xmax),w,h
'''
anno = ET.parse(os.path.join(path))
bbox = list()
for obj in anno.findall('object'):
bndbox_anno = obj.find('bndbox')
# subtract 1 to make pixel indexes 0-based
bbox.append([
int(bndbox_anno.find(tag).text)
for tag in ('ymin', 'xmin', 'ymax', 'xmax')])
size = anno.find('size')
w = size.find('width').text
h = size.find('height').text
bbox = np.stack(bbox).astype(np.float32)
return bbox, w, h
def image_read1(photo_path,hudu=None,tensor=None):
'''
:param photo_path: 图片路径
:param hudu: 是否转化为灰度图
:param tensor: 是否转化为tensor
:return: 图片的数字形式 c,w,h
'''
img = Image.open(photo_path)#将图片读取为图片类
if hudu is True:
img = img.convert("L")#将图片转化为单通道灰度图
img = np.array(img)#将图片转化为数组
if tensor is True:
img = torch.unsqueeze(torch.tensor(img),dim=0)#将numpy转化为tensor
else:
img = np.expand_dims(img,axis=0)
else:
if tensor is True:
img = np.array(img).transpose((2,0,1))#将图片转化为数组,并且将通道数放在第一位
img = torch.tensor(img)#将numpy转化为tensor
else:
img = np.array(img).transpose((2, 0, 1)) # 将图片转化为数组,并且将通道数放在第一位
return img
def all(img_file,xml_file):
'''
Args:
img_file: 图片的文件夹
xml_file: xml文件夹
Returns:
一张一张的显示带有框的图片
'''
img_paths = glob.glob(img_file+"/*.jpg")
for img_path in img_paths:
print(img_path)
img = image_read1(img_path, tensor=True)
if os.path.exists(os.path.join(xml_file,
os.path.basename(img_path).split(".")[0] + ".xml")) is False:
print(os.path.join(xml_file,
os.path.basename(img_path).split(".")[0] + ".xml"),":不存在")
continue
img_xml = os.path.join(xml_file, os.path.basename(img_path).split(".")[0] + ".xml")
bbonx, w_orial, h_orial = box(img_xml)
pali_box(img,bbonx)
time.sleep(10)#设置每次等待时间
if __name__ == '__main__':
all("./VOCdevkit/VOC2007/JPEGImages","./VOCdevkit/VOC2007/Annotations")
每个函数都写明了用途,使用只需要更改all中的参数即可
注意个小细节,由于我用了cv2里的函数来画框并展示,所以此代码中即用了PIL也用了CV2,
注意这两个的读取图片的差别,PIL:RGB;cv2:BGR
还有注意cv2.rectangle函数,传入的图像最好用img.copy(),否则容易报错。