VPGNet数据集解析且将分割标签转化为目标检测标签
VPGNet交通标志数据集
VPGNet是一个包含了四个场景(晴天，夜间，阴雨等）下的交通路况数据集，其主要是针对地面标志（箭头线，人行道）的任务，该数据集一共包括以下17类地面标注
{0:'background', 1:'lane_solid_white',
2:'lane_broken_white', 3:'lane_double_white',
4:'lane_solid_yellow', 5:'lane_broken_yellow',
6:'lane_double_yellow', 7:'lane_broken_blue',
8:'lane_slow', 9:'stop_line',
10:'arrow_left', 11:'arrow_right',
12:'arrow_go_straight', 13:'arrow_u_turn',
14:'speed_bump', 15:'crossWalk',
16:'safety_zone', 17:'other_road_markings',该数据集下载地址为：VPGNet数据集下载
下载解压得到的数据集文件格式为
Scene_x
—%datetime_aaa_xxx
-----id.mat
该数据集的所有有用信息均保存在一个(640×480×5)的mat文件中。
VPGNet数据集内容解析
数据集的关键信息均保存在(6404805)的数组中，前三个通道为图片的rgb三通道信息，第四维通道为图像的17类二值分割图像(像素值为0-17)，第五维通道为该论文提出的vanish point信息(我这里不需要用)。为了利用该数据集，我们需要将其信息分为原始图像images和标签文件labels两部分。python中的scipy库提供了loadmat函数能解析.mat文件。
from scipy.io import loadmat
m=loadmat(mat_file)loadmat函数返回的是字典形式的变量，再继续解析字典m即可分别得到images与labels，完整代码如下：
#取.mat文件，并将其转为图片
import glob
import cv2
import os
import numpy as np
from scipy.io import loadmat
from tqdm import tqdm
root_path='/data/dataset/TrafficMarking/VPGNet/scene_2' #存放mat文件的地址
#save_path='/data/dataset/TrafficMarking/VPGNet/images' #保存原始rgb图像地址
save_path='/data/dataset/TrafficMarking/VPGNet/masks/' #保存原始分割mask图像
datefiles=[os.path.join(root_path,f) for f in os.listdir(root_path)]
pbar=tqdm(datefiles)
for df in pbar: #mat_file='/data/dataset/TrafficMarking/VPGNet/scene_1/20160512_1424_46/000001.mat'
mat_files=glob.glob(df+'/*.mat')
for mat_file in mat_files:
m=loadmat(mat_file) #mat文件的读取格式为dict，按照key,values处理 VPGNet图片信息储存在['rgb_seg_vp']里 （480，640，5）
rgb_img=m['rgb_seg_vp'][:,:,:3]
#修改通道 使之恢复为为原始的rgb图像
b=rgb_img[:,:,0]
g=rgb_img[:,:,1]
r=rgb_img[:,:,2]
rgb_img=cv2.merge((r,g,b))
seg=m['rgb_seg_vp'][:,:,3]*15 #*15为了可视化更清楚
vp=m['rgb_seg_vp'][:,:,4]
file_name=save_path+'/'+root_path.split('/')[-1]+'_'+mat_file.split('/')[-2]+'_'+mat_file.split('/')[-1].split('.')[0]
#保存原始mask图像
#print(np.unique(seg))
cv2.imwrite('%s.png'% file_name,seg)
#保存rgb图像
#file_path='%s.png'%file_name
#cv2.imwrite('%s.png'% file_name,rgb_img.astype(np.uint8))
print( 'finish getting segmentations!')这样即可将图片与分割标签分别保存下来。接下来即按一定比例随机划分训练集与验证集。
import os
import random
import shutil
import numpy as np
from tqdm import tqdm
def split_train_val(filedir):
#随机划分训练集与验证集
val_dir='/data/dataset/TrafficMarking/VPGNet/val/'
np.random.seed(42)
pathdirs=os.listdir(filedir)
ratio=0.2 #验证集比例
number_imgs=len(pathdirs)
val_numbers=int(number_imgs*ratio)
val_samples=random.sample(pathdirs,val_numbers)
for val_sample in val_samples:
shutil.move(filedir+val_sample,val_dir+val_sample)
return
def match_train_val(imgs_dir,labels_dir,saved_path):
#匹配训练集和验证集的标注信息，对应保存
imgs_names=os.listdir(imgs_dir)
anno_names=os.listdir(labels_dir)
pbar=tqdm(imgs_names)
for img_name in pbar:
img_name=img_name.split('.')[0]
anno_name=img_name+'.png' #匹配到的mask标签
#anno_name=img_name+'.txt'
if anno_name in anno_names:
shutil.move(labels_dir+anno_name,saved_path)
return
if __name__=='__main__':
imgs_dir='/data/dataset/TrafficMarking/VPGNet/images/train'
labels_dir='/data/dataset/TrafficMarking/VPGNet/masks/'
saved_path='/data/dataset/TrafficMarking/VPGNet/masks/train'
#filedir='/data/dataset/TrafficMarking/VPGNet/imgs/'
#split_train_val(filedir=filedir)
match_train_val(imgs_dir,labels_dir,saved_path)首先修改filedir为保存数据集图片的地址，执行split_train_val（）将得到验证集图片，剩下的即作为训练集图片。再修改参数执行两次match_train_val()，其中imgs_dir为保存好的训练集，验证集图片的地址，labels_dir为原始所有mask标签的地址，saved_path为保存训练集验证集标签的最终地址。这样即将VPGNet数据集成功分割成为了训练集与验证集两部分，可用于后续的分割模型的训练。其分割标签图片如下所示：
分割标签到检测标签的转化
在我们的使用中，需要将mask部分语义分割标签转化为bbox检测标签，我们主要 通过对于mask 的连通域分析，利用opencv的cv2.connectedComponents()函数来完成该功能，以下提供了完整的代码。
我们选取了提取[10,11,12,13,15,16]几个类别，将其语义分割标签转化为检测框标签。
#将seg标注信息转化为bbox信息
import os
import cv2
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from tqdm import tqdm
from scipy.io import loadmat
import glob
def get_seg(root_path):
#返回原始图片，语义分割标签，文件名称
save_path='/data/dataset/TrafficMarking/VPGNet/masks/' #保存原始分割mask图像
segs=[]
masks=[] #rgb原始图像
file_names=[]
datefiles=[os.path.join(root_path,f) for f in os.listdir(root_path)]
#datefiles=[os.path.join(root_path)]
pbar=tqdm(datefiles)
for df in pbar:
mat_files=glob.glob(df+'/*.mat')
for mat_file in mat_files:
m=loadmat(mat_file) #mat文件的读取格式为dict，按照key,values处理 VPGNet图片信息储存在['rgb_seg_vp']里 （480，640，5） 分割信息保存在seg
rgb_img=m['rgb_seg_vp'][:,:,:3]
#修改通道 使之恢复为为原始的rgb图像
b=rgb_img[:,:,0]
g=rgb_img[:,:,1]
r=rgb_img[:,:,2]
rgb_img=cv2.merge((r,g,b))
seg=m['rgb_seg_vp'][:,:,3]
#file_name=root_path+'/'+mat_file.split('/')[-1].split('.')[0]
file_name=save_path+'/'+root_path.split('/')[-1]+'_'+mat_file.split('/')[-2]+'_'+mat_file.split('/')[-1].split('.')[0]
seg[(seg!=10)&(seg!=11)&(seg!=12)&(seg!=13)&(seg!=15)&(seg!=16)]=0
segs.append(seg)
masks.append(rgb_img)
file_names.append(mat_file)
print('get segmentations finish!')
return masks,segs,file_names
def draw_box(img, boxes, colors,clses):
""" plots one bounding box on image img """
lw = max(round(sum(img.shape) / 2 * 0.003), 2)
tf = max(lw - 1, 1) # font thickness
for box, color,cls in zip(boxes, colors,clses):
cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, thickness=2, lineType=cv2.LINE_AA)
cv2.putText(img,cls , (box[0], box[1] - 2), 0, lw / 3,color, thickness=tf, lineType=cv2.LINE_AA)
#plt.imshow(img)
#plt.axis('off')
#plt.show()
def get_bbox(ori_masks,seg_masks,file_names, needed):
""" get bbox from semantic label """
print('start converting')
#for mask_2D in seg_masks:
for canvas,mask_2D,name in zip(ori_masks,seg_masks,file_names):
#print(canvas,mask_2D)
h=canvas.shape[0]
w=canvas.shape[1]
#保存标注信息的地址
anno_dir='/data/dataset/TrafficMarking/VPGNet/labels'
img_name=name.replace('mat','txt')
#print(img_name.split('/'))
anno_name=root_path.split('/')[-1]+'_'+name.split('/')[-2]+'_'+img_name.split('/')[-1]
anno_path=os.path.join(anno_dir,anno_name.strip())
print(anno_path)
mask_to_save = np.zeros_like(mask_2D)
# instances are encoded as different colors
obj_ids = np.unique(mask_2D)
# split the color-encoded mask into a set of binary masks
masks = mask_2D == obj_ids[:, None, None]
# get bounding box coordinates for each mask
num_objs = len(obj_ids)
boxes, colors,clses = [], [],[]
for i in range(num_objs):
clsid = obj_ids[i]
#name = list(label_dict.keys())[int(id)]
if clsid in needed:
binary = masks[i].astype(np.int8)
num_labels, labels = cv2.connectedComponents(binary, connectivity=8, ltype=cv2.CV_16U)
for id_label in range(1, num_labels):
temp_mask = labels == id_label
pos = np.where(temp_mask)
xmin = np.min(pos[1])+5
xmax = np.max(pos[1])-5
ymin = np.min(pos[0])+5
ymax = np.max(pos[0])-5
#保存为yolo格式的txt信息
cx=np.array((xmin+xmax)/(2*w))
cy=np.array((ymin+ymax)/(2*h))
bw=np.array((xmax-xmin)/w)
bh=np.array((ymax-ymin)/h)
# filter result by setting threshold of width and hegith: 0.04
if bw >= 0.04 and bh >= 0.04:
boxes.append([xmin, ymin, xmax, ymax])
color = [0,255,0]
colors.append(color)
clses.append(label_dict[clsid])
mask_to_save[pos] = id_label
with open(anno_path,'a+') as f:
f.writelines([str(clsid),' ',str(cx),' ',str(cy),' ',str(bw),' ',str(bh),'\n'])
f.close()
# draw bbox
#draw_box(canvas, boxes, colors,clses)
#cv2.imwrite('/home/nianliu/wangxx/yolov5_onnx_inference/bbox.png',canvas
if __name__=='__main__':
needed_clsid=[10,11,12,13,15,16]
root_path='/data/dataset/TrafficMarking/VPGNet/scene_3' #四个scene文件夹分别处理
ori_masks,segs,file_names=get_seg(root_path)
get_bbox(ori_masks,segs,file_names,needed_clsid)
通过上面的代码即可将分割标签转化为yolo格式的目标检测bbox标签。
