MNIST.py 文件源码

def process_digits(all_data, all_labels, data_directory, args):
    """ Generate data and saves in the appropriate format """

    for s in range(len(flags['all_names'])):
        split = flags['all_names'][s]
        print('Processing {0} Data'.format(split))
        key = 'train' if split == 'train' else 'eval'

        # Create writer (tf_records) or Image/Annotations/Names directories (PNGs)
        if args[key] == 'tfrecords':
            tf_writer = tf.python_io.TFRecordWriter(data_directory + 'clutteredMNIST_' + split + '.tfrecords')
        elif args[key] == 'PNG':
            make_Im_An_Na_directories(data_directory)
        else:
            raise ValueError('{0} is not a valid data format option'.format(args[key]))

        # Generate data
        for i in trange(flags['nums'][split]):
            # Generate cluttered MNIST image
            im_dims = [im_dims_generator(), im_dims_generator()]
            num_digits = num_digits_generator()
            img, gt_boxes = gen_nCluttered(all_data[s], all_labels[s], im_dims, num_digits)

            # Save data
            if args[key] == 'tfrecords':
                img = np.float32(img.flatten()).tostring()
                gt_boxes = np.int32(np.array(gt_boxes).flatten()).tostring()
                tf_write(img, gt_boxes, [flags['im_dims'], flags['im_dims']], tf_writer)
            elif args[key] == 'PNG':
                fname = split + '_img' + str(i)
                imsave(data_directory + 'Images/' + fname + '.png', np.float32(img))
                np.savetxt(data_directory + 'Annotations/' + fname + '.txt', np.array(gt_boxes), fmt='%i')
                with open(data_directory + 'Names/' + split + '.txt', 'a') as f:
                    f.write(fname + '\n')


###############################################################################
# Image generation functions
###############################################################################