furg-fire-datase解析,坐标提取,label制作

训练环境: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)


 

  • 1
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值