proc_images.py 文件源码

def main():
    assert os.path.isdir(args.image_path), '%s directory not found' % args.mage_path

    for dataset in ['train', 'test', 'val']:
        out_dir_im = os.path.join(args.output_path, dataset)
        if not os.path.isdir(out_dir_im):
            os.makedirs(out_dir_im)

        out_dir_an = os.path.join(args.output_path, dataset + 'annot')
        if not os.path.isdir(out_dir_an):
            os.makedirs(out_dir_an)

        fid = open(os.path.join(args.output_path, '%s_images.csv' % dataset), 'w')
        # print header
        fid.write('image,labels\n')
        fns = glob(os.path.join(args.image_path, dataset, '*.png'))

        for fn in fns:

            fn_image = os.path.abspath(fn)
            fn_annot = os.path.split(fn_image)
            fn_annot = os.path.join(fn_annot[0] + 'annot', fn_annot[1])

            im = imread(fn_image)
            annot = imread(fn_annot)
            out_size = (256, 512)
            im = imresize(im, out_size)
            annot = imresize(annot, out_size, interp='nearest')


            fn_image_out = os.path.abspath(os.path.join(out_dir_im,
                                                        os.path.basename(fn_image)))
            fn_annot_out = os.path.abspath(os.path.join(out_dir_an,
                                                        os.path.basename(fn_image)))
            imsave(fn_image_out, im)
            imsave(fn_annot_out, annot)

            fid.write('%s,%s\n' %(fn_image_out, fn_annot_out))
        fid.close()