updated eval.py

Author: rkargon

eval.py

@@ -4,68 +4,44 @@
 
 import numpy as np
 import tensorflow as tf
+import time
 
 from model import CNNModel
 from model import restore_model
-from preprocessing import read_object_classes, image_to_np_array, labels_to_np_array, get_patch, save_labels_array
-
-
-def test_model(sess, model, images, labels, patch_size, output_dir=None, category_colors=None):
-    """
-    Tests the model on the given images and labels
-    :param patch_size:
-    :param sess: The tensorflow session within which to run the model
-    :param model: An rCNN model
-    :param images: A series of image filenames
-    :param labels: A series of label filenames (corresponding images/labels should be in the same order)
-    :param output_dir: An (optional) directory in which to store predicted labels as images.
-    :param category_colors: A mapping of category index to color, to create images
-    """
-    for image_f, label_f in zip(images, labels):
-        print "Testing on image %s..." % image_f
-        image = image_to_np_array(image_f, float_cols=True)
-        labels = labels_to_np_array(label_f)
+from preprocessing import read_object_classes, image_to_np_array, labels_to_np_array, get_patch, save_labels_array, \
+    FROM_GAMES, DATASETS
+
+
+def test_model(sess, model, dataset_iter, color_map, output_dir):
+    total_accuracy = 0
+    class_correct_counts = np.zeros(model.num_classes)
+    class_total_counts = np.zeros(model.num_classes)
+    i = 0
+    for image, labels in dataset_iter():
+        i += 1
+        start_time = time.time()
         h, w, _ = image.shape
-        image = image[:h//2, :w//2, :]
-        h, w, _ = image.shape
-        labels = labels[:h, :w]
-        predicted_labels = np.zeros([h, w], dtype=np.uint8)
-        pixels_correct = 0
-        error_for_image = 0
-        i = 0
-
-        for y in range(patch_size, h - patch_size):
-            # # for debug, only do first 10K
-            # if i > 1e4:
-            #     break
-
-            for x in range(patch_size, w - patch_size):
-                i += 1
-                input_image = get_patch(image, (y, x), patch_size)
-                input_image = np.append(input_image,
-                                        np.zeros(shape=[patch_size, patch_size, model.num_classes], dtype=np.float32),
-                                        axis=2)
-                input_label = labels[y, x]
-                feed_dict = {model.inpt: [input_image], model.output: input_label}
-
-                error, logits = sess.run([model.errors[1], model.logits[1]], feed_dict=feed_dict)
-                error_for_image += error
-                output_label = np.argmax(logits)
-                if output_label == input_label:
-                    pixels_correct += 1
-                predicted_labels[y, x] = output_label
-
-                if i % 1000 == 0:
-                    print "%d/%d pixels done..." % (i, (h - 2 * patch_size) * (w - 2 * patch_size))
-
-        # print "Tested on image %s: Accuracy is %.2f%%, error per pixel is %f." % (
-        #     image_f, (100.0 * pixels_correct) / i, error_for_image / i)
-        if output_dir is not None:
-            if category_colors is None:
-                raise ValueError("Color index not provided, can't output images.")
-            output_filename = os.path.join(output_dir, os.path.splitext(os.path.basename(label_f))[0] + '_test.png')
-            print "output: ", output_filename
-            save_labels_array(predicted_labels, output_filename=output_filename, colors=category_colors)
+
+        input_image = np.append(image, np.zeros(shape=[h, w, model.num_classes], dtype=np.float32), axis=2)
+        feed_dict = {model.inpt: [input_image], model.output: [labels]}
+        logits, error = sess.run([model.logits[1], model.loss], feed_dict=feed_dict)
+        predicted_labels = np.argmax(logits[0], axis=2)
+        true_labels = labels[::4, ::4]
+
+        correct_labels = np.equal(predicted_labels, true_labels)
+        accuracy = np.mean(correct_labels)
+        total_accuracy += accuracy
+
+        for c in range(model.num_classes):
+            current_class_labels = np.equal(true_labels, c)
+            class_total_counts[c] += np.sum(current_class_labels)
+            class_correct_counts[c] += np.sum(np.equal(true_labels, c) * correct_labels)
+
+        print "Error: %f Accuracy: %f (time: %.1fs)" % (error, accuracy, time.time() - start_time)
+
+    print "%d Images, Total Accuracy: %f" % (i, total_accuracy/i)
+    print "Per Class accuracy:", class_correct_counts / class_total_counts
+    print np.sum(class_correct_counts / class_total_counts)
 
 
 def main():
@@ -74,8 +50,9 @@ def main():
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument('--model', type=str, help='Filename of saved model')
     parser.add_argument('--category_map', type=str, help='File that maps colors ')
-    parser.add_argument('--images', type=str, nargs='+', help='Filename of test image')
-    parser.add_argument('--labels', type=str, nargs='+', help='Filename of test labels')
+    parser.add_argument('--dataset', type=str, default=FROM_GAMES, choices=DATASETS.keys(),
+                        help='Type of dataset to use. This determines the expected format of the data directory')
+    parser.add_argument('--data_dir', type=str, help='Directory for image and label data')
     parser.add_argument('--output_dir', type=str, default=None,
                         help='Directory to store model output. By default no output is generated.')
     parser.add_argument('--patch_size', type=int, default=67, help='Size of input patches')
@@ -84,14 +61,17 @@ def main():
     # load class labels
     category_colors, category_names, names_to_ids = read_object_classes(args.category_map)
     num_classes = len(category_names)
+
+    # load dataset
+    def dataset_func(): return DATASETS[args.dataset](args.data_dir)
+
     # TODO don't hardcode these (maybe store them in config file?)
     model = CNNModel(25, 50, 1, num_classes, 1e-4, num_layers=2)
 
     sess = tf.Session()
     restore_model(sess, args.model)
 
-    test_model(sess, model, args.images, args.labels, patch_size=args.patch_size, output_dir=args.output_dir,
-               category_colors=category_colors)
+    test_model(sess, model, dataset_func)
 
 
 if __name__ == '__main__':

models/checkpoint

@@ -1,2 +1,2 @@
-model_checkpoint_path: "saved_model.ckpt"
-all_model_checkpoint_paths: "saved_model.ckpt"
+model_checkpoint_path: "stanford_model_10.ckpt"
+all_model_checkpoint_paths: "stanford_model_10.ckpt"

train.py

@@ -14,11 +14,6 @@ def train(sess, model, dataset_iter, num_epochs, patch_size, patches_per_image=1
     for i in range(num_epochs):
         print 'Running epoch %d/%d...' % (i + 1, num_epochs)
         for image, labels in dataset_iter():
-            for row in labels:
-                for l in row:
-                    if l < 0 or l >= model.num_classes:
-                        print "INVALID label:", l
-
             start_time = time.time()
             h, w, _ = image.shape