generate_dset_for_synthtext.py

"""
This script will help you to create dset.h5 file for synthtext data generation.
SynthText: https://github.com/ankush-me/SynthText
First download imnames.cp, bg_img, depth.h5, seg.h5 
Rename imnames.cp to imagesName.txt
Run generate_dset_for_synthtext.py
It will create dset.h5
Then put it inside SynthText/data and run
python gen.py --viz

Reference link: https://blog.csdn.net/baidu_14831657/article/details/77498467

"""

import numpy as np
import h5py
import os, sys, traceback
import os.path as osp
import matplotlib.image as mpimg
 
ImageListDir='imagesName.txt' # just rename imnames.cp to imagesName.txt
ImagesDir='bg_img/'
 
seg_db = h5py.File('seg.h5','r')
depth_db = h5py.File('depth.h5','r')
db = h5py.File('dset.h5','w')  
 
input = open(ImageListDir, 'r')  
imagesName = input.readlines()
 
image=db.create_group("image")
depth=db.create_group("depth")
seg=db.create_group("seg")
num=0
for name in seg_db['mask']:
    name=name.rstrip("\n")
    try:
        img = mpimg.imread(ImagesDir+name)
        d=depth_db[name][:]
        s=seg_db['mask'][name][:]
    except:
        continue
    image[name]=img
    s_max=np.amax(s)
    label=range(1,s_max+1)
    area=range(1,s_max+1)
    for i in label:
        area[i-1]=int(np.sum(s==i))
    depth[name]=d
    seg[name]=s
    seg[name].attrs['area']=area
    seg[name].attrs['label']=label
 
    num=num+1
    print num
    print name + 'is ok!\n'
 
db.close()
seg_db.close()
depth_db.close()