前面我们讲了怎么将各种格式的数据集标签文件进行转换,由于txt格式的标签不同数据集的表示方式可能不一样,将txt格式标注转成voc格式的时候最容易出错。如果voc格式是正确的,转化成coco格式应该也就不会错。
所以我们最好在训练之前将数据集的标注框在原图上打印一下,看看标签有没有出错。
由于很多模型都会用到标准的voc或者coco数据集,所以我们尽量将数据集转成这两种标准格式。下面代码实现将voc数据集的xml标注框打印在对应的原图上并保存在visual_boxes文件夹下。
注意将classes修改成自己的类别,顺序无所谓,后面会根据类别名索引。把原图路径、xml路径和保存标注图片路径填好就可以运行了。
# show_labels_img.py
import os
import os.path
import numpy as np
import xml.etree.ElementTree as xmlET
from PIL import Image, ImageDraw
from tqdm import tqdm
# classes = ('__background__', # always index 0
# 'aeroplane', 'bicycle', 'bird', 'boat',
# 'bottle', 'bus', 'car', 'cat', 'chair',
# 'cow', 'diningtable', 'dog', 'horse',
# 'motorbike', 'person', 'pottedplant',
# 'sheep', 'sofa', 'train', 'tvmonitor')
classes=('billboard',) # 如果不加后面的',',输出的标签只有一个字符b
# 这里我的图片和xml文件放在同一文件夹了,所以后面用了一个if语句判断是否为标签
file_path_img = 'F:/Billboard/mydata/images'
file_path_xml = 'F:/Billboard/mydata/images'
save_file_path = 'F:/Billboard/mydata/visual_boxes'
pathDir = os.listdir(file_path_xml)
#for idx in range(10): # 图片过多可以选择只看几张,如果图片跟标签放一起,剔除非标签文件后只输出5张标注图片
for idx in tqdm(range(len(pathDir))):
filename = pathDir[idx]
if filename[-3:]=='xml': # 防止图片跟标签放一起,读取的时候出错,将xml单独放在一个文件夹时可以去掉判断
tree = xmlET.parse(os.path.join(file_path_xml, filename))
objs = tree.findall('object')
num_objs = len(objs)
boxes = np.zeros((num_objs, 5), dtype=np.uint16)
for ix, obj in enumerate(objs):
bbox = obj.find('bndbox')
# Make pixel indexes 0-based
x1 = float(bbox.find('xmin').text) - 1
y1 = float(bbox.find('ymin').text) - 1
x2 = float(bbox.find('xmax').text) - 1
y2 = float(bbox.find('ymax').text) - 1
cla = obj.find('name').text
label = classes.index(cla)
boxes[ix, 0:4] = [x1, y1, x2, y2]
boxes[ix, 4] = label
image_name = os.path.splitext(filename)[0]
img = Image.open(os.path.join(file_path_img, image_name + '.jpg'))
draw = ImageDraw.Draw(img)
for ix in range(len(boxes)):
xmin = int(boxes[ix, 0])
ymin = int(boxes[ix, 1])
xmax = int(boxes[ix, 2])
ymax = int(boxes[ix, 3])
draw.rectangle([xmin, ymin, xmax, ymax], outline=(0, 255, 0))
draw.text([xmin, ymin-10], classes[boxes[ix, 4]], (0, 255, 0))
img.save(os.path.join(save_file_path, image_name + '.png'))
如果是根据load数据部分的代码改的txt标签,可以根据实际情况进行打印。
以上面的格式为例,0表示第0类即’billboard’,后面四位分别表示x1,x2,y1,y2, 而且是用长和宽归一化后的值。
将图片00013ec42a6a93a5.jpg
的标注框打印在原图上,代码实现如下:
import cv2
def show_labels_img(imgname):
img = cv2.imread(DATASET_PATH + imgname + ".jpg")
# img = cv2.imdecode(np.fromfile(DATASET_PATH + imgname + ".jpg", dtype=np.uint8), flags=cv2.IMREAD_COLOR) # 中文路径使用
h, w = img.shape[:2]
print(w,h)
label = []
with open("/home/guest-1/dwc/Billboard/dataset/labels/train/"+imgname+".txt",'r') as flabel:
for label in flabel:
label = label.split(' ')
label = [float(x.strip()) for x in label]
print(CLASSES[int(label[0])])
pt1 = (int(label[1] * w-label[3] * w/2), int(label[2] * h-label[4] * h/2))
pt2 = (int(label[1] * w+label[3] * w/2), int(label[2] * h+label[4] * h/2))
cv2.putText(img,CLASSES[int(label[0])],pt1,cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,255))
cv2.rectangle(img,pt1,pt2,(0,0,255,2))
cv2.imwrite("./show_label_img.jpg",img)
if __name__ == '__main__':
# CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
# 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
CLASSES=['Billboard']
DATASET_PATH='/home/guest-1/dwc/Billboard/dataset/images/train/'
show_labels_img('00013ec42a6a93a5')
需要注意的是,如果用cv2读取图片,路径不要有中文字符,否则会img会返归None type。如果确实需要用中文路径,把读取图片部分的代码换成img = cv2.imdecode(np.fromfile(中文路径, dtype=np.uint8), flags=cv2.IMREAD_COLOR)
,保存图片部分代码相应改成cv2.imencode('.jpg', img)[1].tofile('中文路径')
即可。