extract bdd images for hfnet training

bdd_extract_images_for_hfnet.py

#!/usr/bin/env python

import json
import tqdm
from pathlib import Path
import shutil

def load_night_and_dawn_list(fn):
    print('loading night&dawn image list from', fn)
    fns_night = []
    fns_dawn = []
    with open(fn, 'r') as f:
        db = json.load(f)
    for item in db:
        name = item['name']
        tod = item['attributes']['timeofday']
        if tod.find('night')>=0:
            fns_night.append(name)
        elif tod.find('dawn')>=0:
            fns_dawn.append(name)
    print('night:%d dawn:%d'%(len(fns_night), len(fns_dawn)))
    return fns_night, fns_dawn

def copy_images(fp_src:Path, fp_dst:Path, fns:list):
    print("copying images from %s to %s"%(fp_src.name, fp_dst.name))
    fp_dst.mkdir(parents=True, exist_ok=True)
    for fn in tqdm.tqdm(fns):
        shutil.copy(fp_src/fn, fp_dst/fn);

def move_images(fp_src:Path, fp_dst:Path, fns:list):
    print("moving images from %s to %s"%(fp_src.name, fp_dst.name))
    fp_dst.mkdir(parents=True, exist_ok=True)
    for fn in tqdm.tqdm(fns):
        shutil.move(fp_src/fn, fp_dst/fn);

def main(work_dir:Path):
    # pre-defined dirs
    work_dir = Path(work_dir)
    fp_label = work_dir/'labels'
    fp_image = work_dir/'images/100k'
    fp_dst = work_dir/'bdd'

    # sub dirs
    fn_train = fp_label/'bdd100k_labels_images_train.json'
    fn_val = fp_label/'bdd100k_labels_images_val.json'
    fp_image_train = fp_image/'train'
    fp_image_val = fp_image/'val'
    fp_dst_night = fp_dst/'night_images_vga'
    fp_dst_dawn = fp_dst/'dawn_images_vga'

    # get image lists
    fns_night_train, fns_dawn_train = load_night_and_dawn_list(fn_train)
    fns_night_val, fns_dawn_val = load_night_and_dawn_list(fn_val)
    print('total image num:', len(fns_night_train)+len(fns_dawn_train)
        +len(fns_night_val)+len(fns_dawn_val))

    # copy images
    move_images(fp_image_train, fp_dst_night, fns_night_train)
    move_images(fp_image_train, fp_dst_dawn, fns_dawn_train)
    move_images(fp_image_val, fp_dst_night, fns_night_val)
    move_images(fp_image_val, fp_dst_dawn, fns_dawn_val)

if __name__=='__main__':
    main('.')