util.py 1.9 KB
Newer Older
HypoX64's avatar
preview  
HypoX64 已提交
1
import os
H
hypox64 已提交
2
import shutil
H
hypox64 已提交
3

HypoX64's avatar
preview  
HypoX64 已提交
4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
def Traversal(filedir):
    file_list=[]
    for root,dirs,files in os.walk(filedir): 
        for file in files:
            file_list.append(os.path.join(root,file)) 
        for dir in dirs:
            Traversal(dir)
    return file_list

def is_img(path):
    ext = os.path.splitext(path)[1]
    ext = ext.lower()
    if ext in ['.jpg','.png','.jpeg','.bmp']:
        return True
    else:
        return False

def is_video(path):
    ext = os.path.splitext(path)[1]
    ext = ext.lower()
H
hypox64 已提交
24
    if ext in ['.mp4','.flv','.avi','.mov','.mkv','.wmv','.rmvb','.mts']:
HypoX64's avatar
preview  
HypoX64 已提交
25 26 27
        return True
    else:
        return False
H
hypox64 已提交
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42

def is_imgs(paths):
    tmp = []
    for path in paths:
        if is_img(path):
            tmp.append(path)
    return tmp

def is_videos(paths):
    tmp = []
    for path in paths:
        if is_video(path):
            tmp.append(path)
    return tmp  

HypoX64's avatar
preview  
HypoX64 已提交
43 44 45
def  writelog(path,log):
    f = open(path,'a+')
    f.write(log+'\n')
H
hypox64 已提交
46
    f.close()
HypoX64's avatar
preview  
HypoX64 已提交
47

H
hypox64 已提交
48 49 50 51 52 53 54
def makedirs(path):
    if os.path.isdir(path):
        print(path,'existed')
    else:
        os.makedirs(path)
        print('makedir:',path)

H
hypox64 已提交
55
def clean_tempfiles(tmp_init=True):
HypoX64's avatar
preview  
HypoX64 已提交
56
    if os.path.isdir('./tmp'):   
H
hypox64 已提交
57 58 59 60 61 62 63
        shutil.rmtree('./tmp')
    if tmp_init:
        os.makedirs('./tmp')
        os.makedirs('./tmp/video2image')
        os.makedirs('./tmp/addmosaic_image')
        os.makedirs('./tmp/mosaic_crop')
        os.makedirs('./tmp/replace_mosaic')
H
hypox64 已提交
64
        os.makedirs('./tmp/mosaic_mask')
H
hypox64 已提交
65
        os.makedirs('./tmp/ROI_mask')
H
hypox64 已提交
66
        os.makedirs('./tmp/ROI_mask_check')
H
hypox64 已提交
67 68 69 70 71 72

def file_init(opt):
    if not os.path.isdir(opt.result_dir):
        os.makedirs(opt.result_dir)
        print('makedir:',opt.result_dir)
    clean_tempfiles()
H
HypoX64 已提交
73 74 75 76 77 78 79 80 81 82

def get_bar(percent,num = 25):
    bar = '['
    for i in range(num):
        if i < round(percent/(100/num)):
            bar += '#'
        else:
            bar += '-'
    bar += ']'
    return bar+' '+str(round(percent,2))+'%'