generate_norb_small.py 文件源码

def process_file(params):
    index, data, base_filename, db_name, C, aug_data = params
    label = index % NUM_CLASSES
    if C==1:
        orig_im = data[0,:,:]
        im = ndimage.interpolation.zoom(orig_im, DOWNSCALE_FACTOR)
    elif C==2:
        im = np.zeros((int(MAT_SHAPE[2]*DOWNSCALE_FACTOR),int(MAT_SHAPE[3]*DOWNSCALE_FACTOR),3))
        orig_im = np.zeros((MAT_SHAPE[2],MAT_SHAPE[3],3))
        im[:,:,0] =  ndimage.interpolation.zoom(data[0,:,:], DOWNSCALE_FACTOR)
        im[:,:,1] =  ndimage.interpolation.zoom(data[1,:,:], DOWNSCALE_FACTOR)
        orig_im[:,:,0] =  data[0,:,:]
        orig_im[:,:,1] =  data[1,:,:]
    else:
        print "Error in reading data to db- number of channels must be 1 or 2"
    im_name = '%s_%d%s' % (base_filename, index,IM_FORMAT)
    scipy.misc.toimage(im, cmin=0.0, cmax=255.0).save(os.path.join(db_name,im_name))
    im_names = [im_name]
    if aug_data:
        degrees = [-20, -10, 10, 20]
        crop_dims = [2, 4, 6, 8]
        for i, degree in enumerate(degrees):
            im_name = '%s_%d_%d%s' % (base_filename,index,degree,IM_FORMAT)
            im_names.append(im_name)
            rot_im = rotate_im(orig_im, degree)
            scipy.misc.toimage(rot_im, cmin=0.0, cmax=255.0).save(os.path.join(db_name,im_name))
        for i, crop_dim in enumerate(crop_dims):
            im_name = '%s_%d_%d%s' % (base_filename,index,crop_dim,IM_FORMAT)
            im_names.append(im_name)
            cr_im = crop_and_rescale(orig_im, crop_dim)        
            scipy.misc.toimage(cr_im, cmin=0.0, cmax=255.0).save(os.path.join(db_name,im_name))
    return label, im_names