MCCNN/train.py at master · jjxxmiin/MCCNN · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
from utils import *
import tensorflow as tf
import mcnn as model
from ops import mse
from tensorflow.python.framework import graph_io

import matplotlib.image as mpimg
import scipy.io as sio

dataset = 'A'

log_dir = "logs"
checkpoint_dir = "checkpoint"
sample_dir = "sample"

learning_rate = 1e-6
epoch = 200000

image = tf.placeholder(tf.float32,shape=[None,None,None,3])
ground_truth = tf.placeholder(tf.float32,shape=[None,None,None,1])


# model
estimate = model.multi_column_cnn(image)


# MSE
loss = mse(estimate,ground_truth)
# SGD : batch size 1
optimizer = tf.train.AdamOptimizer(learning_rate)

train = optimizer.minimize(loss)

tf.summary.scalar("MSE",loss)
summary = tf.summary.merge_all()

saver = tf.train.Saver()

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)

    # tensorboard
    writer = tf.summary.FileWriter(log_dir, sess.graph)

    # model loading
    could_load, checkpoint_counter = load(checkpoint_dir, sess, saver)

    train_image_list, train_gt_list, iteration = get_data_list(dataset, mode='train')

    if could_load:
        counter = checkpoint_counter
        start_epoch = checkpoint_counter // iteration
        print(" [*] Load SUCCESS")
    else:
        counter = 1
        start_epoch = 0
        print(" [!] Load failed...")

    for e in range(start_epoch,epoch):

        for i in range(iteration):
            img,gt_dmp,gt_count = read_train_data(train_image_list[i],train_gt_list[i],scale=4)
            img = input_normalization(img)

            _,prediction,cost,summary_str = sess.run([train,estimate,loss,summary],feed_dict={image : img ,
                                                       ground_truth : gt_dmp
                                                       })


            writer.add_summary(summary_str, counter)
            counter += 1

            print("[{} / {}] [{} / {}] LOSS : {} pre : {} gt : {}".format(epoch,e,iteration,i,cost,prediction.sum(),gt_count))

            # frozen = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ["mcnn/dmp_conv1/Relu"])

            # graph_io.write_graph(frozen, './', 'graph.pb', as_text=False)

        if e % 101 == 0:
            absolute_error = 0.0
            square_error = 0.0

            test_image_list, test_gt_list, total_test_count = get_data_list(dataset, mode='test')

            # validating
            for j in range(total_test_count):

                img, gt_dmp, gt_count = read_test_data(test_image_list[j], test_gt_list[j], scale=4)

                img = input_normalization(img)

                _, prediction, cost, summary_str = sess.run([train, estimate, loss, summary], feed_dict={image: img,
                                                                                                         ground_truth: gt_dmp
                                                                                                         })

                absolute_error = absolute_error + np.abs(np.subtract(gt_count, prediction.sum())).mean()
                square_error = square_error + np.power(np.subtract(gt_count, prediction.sum()), 2).mean()

                print("[{} / {}] [{} / {}] LOSS : {} pre : {} gt : {}".format(epoch, e, total_test_count, j, cost,
                                                                              prediction.sum(), gt_count))

            mae = absolute_error / total_test_count
            rmse = np.sqrt(absolute_error / total_test_count)

            print(str('MAE_' + str(mae) + '_MSE_' + str(rmse)))

        if e % 100 == 0:
            save(checkpoint_dir, "mcnn", counter, sess, saver)
            # graph
            tf.train.write_graph(sess.graph_def, '.', 'graph.pbtxt')

        train_image_list, train_gt_list, _ = get_data_list(dataset,mode='train')