Пример #1
0
t_filenames = tf.constant(t_train_name)

dataset = tf.data.Dataset.from_tensor_slices((x_filenames, t_filenames))
dataset = dataset.map(lambda x, y: _parse_function(x, y, (img_height, img_width), down_scale))
dataset = dataset.batch(batch_size).repeat(1)
iterator = dataset.make_initializable_iterator()
next_batch = iterator.get_next()

x_batch, t_batch = next_batch # get the tf variable of input and target images



unet = UNet(x=x_batch, t=t_batch,
            LR=1e-8, input_shape=[None, img_height, img_width, 3], 
            output_shape=[None, img_height, img_width, class_num], )
unet.optimize(entropy_loss)

sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(iterator.initializer)



saver = tf.train.Saver(max_to_keep=epoch)
saver.restore(sess, './Models/U-Net/unet-'+epoch+'/unet.ckpt')
for ep in range(1):
    total_loss = 0
    counter = 0
    start = time.time()
    for _ in range(int(math.ceil(data_size/batch_size))):
        _, loss = sess.run([unet.training, unet.loss])
    img_height, img_width), down_scale, class_num))
dataset = dataset.shuffle(buffer_size=32).batch(batch_size).repeat(epoch + 1)
iterator = dataset.make_initializable_iterator()
next_batch = iterator.get_next()

x_batch, t_batch = next_batch  # get the tf variable of input and target images

if model_name.lower() == 'unet' or model_name.lower == 'u-net':
    segnet = UNet(
        x=x_batch,
        t=t_batch,
        LR=LR,
        input_shape=[None, img_height, img_width, 3],
        output_shape=[None, img_height, img_width, class_num],
    )
    segnet.optimize(loss_function)
elif model_name.lower() == 'fcn':
    segnet = FCN(
        x=x_batch,
        t=t_batch,
        LR=LR,
        input_shape=[None, img_height, img_width, 3],
        output_shape=[None, img_height, img_width, class_num],
    )
    segnet.optimize(loss_function)
elif model_name.lower() == 'resnet50' or model_name.lower() == 'resnet':
    segnet = FCN_ResNet50(
        x=x_batch,
        t=t_batch,
        LR=LR,
        input_shape=[None, img_height, img_width, 3],