Ejemplo n.º 1
0
def data_testsave(input_dir, output_dir):
    filename = 'TestSet.tfrecord'
    writer = tf.python_io.TFRecordWriter(os.path.join(output_dir + filename))
    for folder in os.listdir(input_dir):
        print('****************' + folder + '***********************')
        img_path = input_dir + '/' + folder + '/Image/'
        lbl_path = input_dir + '/' + folder + '/Label/'
        for imgname in os.listdir(img_path):
            lblname = imgname[:-3] + 'bmp'
            print(imgname + "," + lblname)
            img = cv2.imread(img_path + imgname)
            lbl = cv2.imread(lbl_path + lblname)
            #imgs =  np.dot(img[...,:3], [0.299, 0.587, 0.114])
            imgs = pr.Intensity(img)
            lbls = np.dot(lbl[..., :3], [0.299, 0.587, 0.114]) // 255
            lbls = lbls.astype('uint8')

            img_batch = pr.down_sample(imgs).astype('float64')
            lbl_batch = pr.down_sample(lbls)
            img_batch_norm = pr.contrast_normalization(img_batch)

            example = tf.train.Example(features=tf.train.Features(
                feature={
                    'img_raw': _bytes_feature(img_batch_norm.tostring()),
                    'gt_raw': _bytes_feature(lbl_batch.tostring())
                }))
            writer.write(example.SerializeToString())
    writer.close()
Ejemplo n.º 2
0
def data_trainsave(input_dir, output_dir):

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    filename = 'TrainSet.tfrecord'
    Trans = ['False', 'Mirror', 'Rot90', 'Rot-90', 'Displa', 'Rotate']
    dx = 0
    dy = 0
    theta = 0
    writer = tf.python_io.TFRecordWriter(os.path.join(output_dir + filename))
    for itr in range(2):
        if itr > 6:
            trans_mode = Trans[5]
        elif itr > 4 and itr <= 6:
            trans_mode = Trans[4]
        elif itr <= 4:
            trans_mode = Trans[itr]
        print('____________________' + trans_mode +
              '___________________________')
        for folder in os.listdir(input_dir):
            #img_path = folder+'/Image/'
            #lbl_path = folder+'/Label/'
            for sub_folder in os.listdir(input_dir + '/' + folder + '/'):
                #savenum=0
                #savename = filename+str(savenum)+'.tfrecord'
                #writer = tf.python_io.TFRecordWriter(os.path.join(output_dir+savename))
                img_folder = input_dir + '/' + folder + '/' + sub_folder + '/Image/'
                lbl_folder = input_dir + '/' + folder + '/' + sub_folder + '/Label/'
                print('****************' + sub_folder +
                      '***********************')
                for imgname in os.listdir(img_folder):
                    lblname = imgname[:-3] + 'bmp'
                    print(imgname + "," + lblname)
                    img = cv2.imread(img_folder + imgname)
                    lbl = cv2.imread(lbl_folder + lblname)
                    #imgs =  np.dot(img[...,:3], [0.299, 0.587, 0.114])
                    imgs = pr.Intensity(img)
                    lbls = np.dot(lbl[..., :3], [0.299, 0.587, 0.114]) // 255
                    lbls = lbls.astype('uint8')

                    img_batch = pr.down_sample(imgs).astype('float64')
                    lbl_batch = pr.down_sample(lbls)

                    img_batch_norm = pr.contrast_normalization(img_batch)
                    if trans_mode == 'Displa':
                        dx = random.randint(-40, 40)
                        dy = random.randint(-40, 40)
                        print('dx:' + str(dx) + '...dy:' + str(dy))
                    if trans_mode == 'Rotate':
                        theta = random.randint(-50, 50)
                        print('theta:' + str(dx))
                    img_all_batch = pr.data_argument(img_batch_norm,
                                                     trans_mode, dx, dy, theta)
                    lbl_all_batch = pr.data_argument(lbl_batch, trans_mode, dx,
                                                     dy, theta)

                    img_all = img_all_batch.reshape([1, HEIGHT, WIDTH, 1])
                    lbl_all = lbl_all_batch.reshape([1, HEIGHT, WIDTH, 1])

                    img_sp = {}
                    lbl_sp = {}
                    img_sp = pr.image_splite(img_all, COLS, ROWS)
                    lbl_sp = pr.image_splite(lbl_all, COLS, ROWS)

                    for col in range(COLS):
                        for row in range(ROWS):
                            print('~~~~~~~~' + str(col) + '~~~' + str(row) +
                                  '~~~~~~~~')
                            img_save_batch = img_sp[col][row]
                            lbl_save_batch = lbl_sp[col][row]

                            img_save = img_save_batch.reshape(
                                [SP_HEIGHT, SP_WIDTH, 1])
                            lbl_save = lbl_save_batch.reshape(
                                [SP_HEIGHT, SP_WIDTH, 1]).astype('uint8')

                            example = tf.train.Example(
                                features=tf.train.Features(
                                    feature={
                                        'img_raw':
                                        _bytes_feature(img_save.tostring()),
                                        'gt_raw':
                                        _bytes_feature(lbl_save.tostring())
                                    }))
                            writer.write(example.SerializeToString())
    writer.close()