def train(net, img_ph, ann_ph, title):
# Load data
(train_img, train_ann), (test_img, test_ann) = ear_pen.load_data()
train_ann, _map = to_categorical_4d(train_ann)
print('map: ', _map)
# Train
saver = tf.train.Saver()
loss_list = []
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for i in range(epoch):
loss_sum = 0
for j in range(math.ceil(len(train_img) / batch_size) + 1):
if j == math.ceil(len(train_img) / batch_size):
feed_dict = {
# avoid "could not convert BatchDescriptor" cudnn error
img_ph:
train_img[len(train_img) - batch_size:len(train_img)],
ann_ph:
train_ann[len(train_img) - batch_size:len(train_img)]
}
else:
feed_dict = {
img_ph:
train_img[j * batch_size:j * batch_size + batch_size],
ann_ph:
train_ann[j * batch_size:j * batch_size + batch_size]
}
_loss, _, _pred = sess.run(
[net.loss, net.train_op, net.prediction],
feed_dict=feed_dict)
loss_sum += _loss
_img = formStoredImg(train_img, train_ann, _map, _pred)
if i % save_period == 0:
imageio.imsave(str(i) + '.png', _img)
loss_list.append(loss_sum)
print('iter: ', i, '\tloss: ', loss_sum)
# Store train result
model_name = title[8:]
saver.save(sess,
model_store_path + model_name + '/' + model_name + '.ckpt')
loss_list[0] = loss_list[-1]
statistic(range(int(epoch / save_period)), np.log(loss_list), title)
def work(img_ph, ann_ph, net, model_path):
# Load data
(train_img, train_ann), (test_img, test_ann) = ear_pen.load_data()
test_ann, _map = to_categorical_4d(test_ann)
# Test
saver = tf.train.Saver()
loss_list = []
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess, model_path)
loss_sum = 0
for i in range(math.ceil(len(test_img) / batch_size)):
feed_dict = {
img_ph: test_img[i * batch_size:i * batch_size + batch_size],
ann_ph: test_ann[i * batch_size:i * batch_size + batch_size]
}
_loss, _pred = sess.run([net.loss, net.prediction],
feed_dict=feed_dict)
loss_sum += _loss
loss_list.append(loss_sum)
print('test loss: ', loss_sum)
showImg(test_img, test_ann, _map, _pred)
from utils import to_categorical_4d, to_categorical_4d_reverse
import ear_pen
import numpy as np
"""
This main.py shows how to use the Ear-pen images in a simple way.
Just like keras style, you can call `load_data` function to load data
The returned data are two tuple with image and annotation
Moreover, I provide `to_categorical_4d` function which is just like to_categorical in keras
The different is that it will transfer the annotation as 4D tensor which channel number is the number of classes
To be notice: the annotation should be normalized if you want to convert as one-hot format
"""
if __name__ == '__main__':
# Load data and normalize
(train_x, train_y), (test_x, test_y) = ear_pen.load_data()
# to_categorical
print('to_categorical_4d')
train_y_, _map = to_categorical_4d(train_y)
# to_categorical reverse
print('to_categorical_4d_reverse')
train_y_ = to_categorical_4d_reverse(train_y_, _map)
# Check if there're the same
print('origin shape: ', np.shape(train_y), '\tshape after reversed: ',
np.shape(train_y_))
print('Reversed result is the same? \t', np.array_equal(train_y, train_y_))
from utils import to_categorical_4d, to_categorical_4d_reverse
from FCN import FCN8
from draw import drawRec
import tensorflow as tf
import numpy as np
import ear_pen
import math
batch_size = 1
model_store_path = '../model/FCN-8/FCN-8.ckpt'
video_name = 'move.mp4'
if __name__ == '__main__':
img_ph = tf.placeholder(tf.float32, [None, 104, 78, 3])
ann_ph = tf.placeholder(tf.int32, [None, 104, 78, 1])
net = FCN8(img_ph, ann_ph)
# Load data
(train_img, train_ann), (test_img, test_ann) = ear_pen.load_data()
train_ann = np.asarray(train_ann) / 255
train_ann, _map = to_categorical_4d(train_ann)
# Work
saver = tf.train.Saver()
loss_list = []
with tf.Session() as sess:
# Recover model
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess, model_store_path)
print('restore done...')