训练环境:darknet+yolov3,坐标写入格式:x,y,w,h,如果需要去他格式的坐标可以修改
def convert(size, box)函数。
furg-fire-datase:
火焰检测数据机,涵盖了大部分交通事故中火焰检测的场景。
数据集制作脚本
# -*- coding: utf-8 -*-
# @Time : 2020/7/14 13:23
# @Author : Libin
# @File : ReadXml.py
# @Software: PyCharm
import os
import glob
import shutil
import cv2
debug=False
train_dir='/mnt/temp/'
train_data_path='/mnt/temp/train_set.txt'
path = "./"
# save_path = "C:/Users/91324/Desktop/model/Xml2Txt/"
xml_file = os.listdir(path)
def convert(size, box):
"""
:param size:image size
:param box: minx,miny,maxx,maxy
:return:
"""
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[2])/2.0
y = (box[1] + box[3])/2.0
w = box[2] - box[0]
h = box[3] - box[1]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def rename(base_path):
"""
:param base_path:dataset directory
:return: None
"""
print("***********************rename XML***********************\n")
try:
xmllist=glob.glob(os.path.join(base_path,"*.xml"))
except Exception as e:
print(" %s " %e)
return False
for xml_file in xmllist:
txt_file=xml_file[:-4]+'.txt'
# os.rename(xml_file, txt_file)
shutil.copy(xml_file, txt_file)
read_coordinate(txt_file)
if len(xmllist)==0:
txtlist = glob.glob(os.path.join(base_path, "*.txt"))
for txt_file in txtlist:
read_coordinate(txt_file)
def draw_rect(txt_path,coor):
"""
:param txt_path: label file save to txt_path
:param coor: region coor
:return: None
"""
img_name=txt_path[:-4]+'.jpg'
if not os.path.exists(img_name):
return
img=cv2.imread(img_name)
cv2.rectangle(img,(coor[0],coor[1]),(coor[2],coor[3]),color=(0,255,0),thickness=1)
cv2.imwrite(img_name,img)
def read_coordinate(txt_file):
"""
:param txt_file: read text from txt_file (coordinate )
:return:
"""
global iw,ih
iw=-1
ih=-1
try:
new_frames = False
with open(txt_file,'r') as f:
lines=f.readline()
while lines !="":
print(lines)
"""
Get image size
"""
if lines.strip().find("</frameWidth>") !=-1:
iw=int(lines.strip().split('</frameWidth>')[0].split('<frameWidth>')[1])
elif lines.strip().find("</frameHeight>") !=-1:
ih=int(lines.strip().split('</frameHeight>')[0].split('<frameHeight>')[1])
target_single_frame=[]
# single_box=[]
#检测到了视频帧
if lines.find('frameNumber') !=-1 :
frames_num=int(lines.strip().split('<frameNumber>')[1].split('</frameNumber>')[0]) # <_><frameNumber>0</frameNumber>
print("frames number :%d " %frames_num)
while True:
lines=f.readline()
if lines.find('frameNumber') ==-1:
if lines.find('</_></annotations></_>')!=-1:#包含坐标值的行
need_value=lines.strip().split('</_></annotations></_>')[0]
if need_value=='</_></annotations></_>':
continue
else:
print("read from txt: ",need_value)
minx,miny,w,h=[int(i) for i in need_value.split(' ')]
target_single_frame.append([minx,miny,minx+w,miny+h])
else:
new_frames = True
if len(target_single_frame) !=0:
# label_file=save_path+str(frames_num)+'_train.txt'
label_dir=path+os.path.basename(txt_file)[:-4]
if not os.path.exists(label_dir):
os.makedirs(label_dir, exist_ok=True)
label_file=label_dir+'/'+str(frames_num)+'.txt'
print(label_file)
fw=open(label_file,'w')
for box in target_single_frame:
if debug:
draw_rect(label_file, box)
con_box = convert((iw, ih), box)
fw.write(str('2') + " " + " ".join([str(a) for a in con_box]) + '\n')
fw.close()
if os.path.exists(label_file[:-4] + '.jpg'):
base_label_path = os.path.basename(label_file)
src_img_path = label_file[:-4] + '.jpg'
save_label_to_path = train_dir + base_label_path
save_img_to_path = train_dir + base_label_path[:-4] + '.jpg'
shutil.copy(src_img_path, save_img_to_path)
shutil.copy(label_file, save_label_to_path)
with open(train_data_path, 'a') as tfp:
tfp.write(save_img_to_path + '\n')
break
if lines.strip()=='':
"""写入视频最后一帧数据集"""
# label_file = save_path + str(frames_num) + '_train.txt'
label_file = path + os.path.basename(txt_file)[:-4] + '/' + str(frames_num) + '.txt'
print(label_file)
fw = open(label_file, 'w')
for box in target_single_frame:
if debug:
draw_rect(label_file, box)
con_box = convert((iw, ih), box)
fw.write(str('2') + " " + " ".join([str(a) for a in con_box]) + '\n')
fw.close()
if os.path.exists(label_file[:-4] + '.jpg'):
base_label_path = os.path.basename(label_file)
src_img_path = label_file[:-4] + '.jpg'
save_label_to_path = train_dir + base_label_path
save_img_to_path = train_dir + base_label_path[:-4] + '.jpg'
shutil.copy(src_img_path, save_img_to_path)
shutil.copy(label_file, save_label_to_path)
with open(train_data_path, 'a') as tfp:
tfp.write(save_img_to_path + '\n')
return #文件读完
if new_frames==True:
new_frames=False
continue
else:
lines = f.readline().strip()
except Exception as e:
print(" %s " %e)
return False
def save_video_frame():
"""
:return:
"""
videos = glob.glob(os.path.join(path,"*.mp4"))
print("video path:",videos)
for video_name in videos:
folder_name=path+os.path.basename(video_name)[:-4]
os.makedirs(folder_name, exist_ok=True)
print("="*10+'开始处理视频:%s'%video_name+"="*10)
vc = cv2.VideoCapture(video_name)
c = 0
rval = vc.isOpened()
while rval:
rval, frame = vc.read()
pic_path = folder_name + '/'
if rval:
cv2.imwrite(pic_path + str(c) + '.jpg', frame)
cv2.waitKey(10)
c += 1
else:
break
print('='*10+'%s图片数量:%i'%(video_name,c)+'='*10)
vc.release()
print('save_success')
print(folder_name)
if __name__=='__main__':
# save_video_frame()
rename(path)