训练环境: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 osimport glob
import shutil
import cv2debug=Falsetrain_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.0y = (box[1] + box[3])/2.0w = box[2] - box[0]h = box[3] - box[1]x = x*dww = w*dwy = y*dhh = h*dhreturn (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 Falsefor 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):returnimg=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,ihiw=-1ih=-1try:new_frames = Falsewith 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></_>':continueelse: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 = Trueif 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_pathsave_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')breakif 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_pathsave_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=Falsecontinueelse:lines = f.readline().strip()except Exception as e:print(" %s " %e)return Falsedef 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 = 0rval = 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 += 1else:breakprint('='*10+'%s图片数量:%i'%(video_name,c)+'='*10)vc.release()print('save_success')print(folder_name)if __name__=='__main__':# save_video_frame()rename(path)