fix tf 1.7 test.py support

Author: sthalles

README.md

@@ -42,6 +42,8 @@ Check out the *train.py* file for more input argument options. Each run produces
 
 To evaluate the model, run the *test.py* file passing to it the *model_id* parameter (the name of the folder created inside *tboard_logs* during training).
 
+Note: Make sure the `test.tfrecords` is downloaded and placed inside `./dataset/tfrecords`.
+
 ```
 python test.py --model_id=16645
 ```

preprocessing/read_data.py

@@ -4,7 +4,7 @@
 import tarfile
 import os
 
-def random_flip_image_and_annotation(image_tensor, annotation_tensor):
+def random_flip_image_and_annotation(image_tensor, annotation_tensor, image_shape):
     """Accepts image tensor and annotation tensor and returns randomly flipped tensors of both.
     The function performs random flip of image and annotation tensors with probability of 1/2
     The flip is performed or not performed for image and annotation consistently, so that
@@ -44,10 +44,10 @@ def random_flip_image_and_annotation(image_tensor, annotation_tensor):
                                                         true_fn=lambda: tf.image.flip_left_right(annotation_tensor),
                                                         false_fn=lambda: annotation_tensor)
 
-    return randomly_flipped_img, tf.reshape(randomly_flipped_annotation, original_shape)
+    return randomly_flipped_img, tf.reshape(randomly_flipped_annotation, original_shape, name="reshape_random_flip_image_and_annotation"), image_shape
 
 
-def rescale_image_and_annotation_by_factor(image, annotation, nin_scale=0.5, max_scale=2):
+def rescale_image_and_annotation_by_factor(image, annotation, image_shape, nin_scale=0.5, max_scale=2):
     #We apply data augmentation by randomly scaling theinput images(from 0.5 to 2.0)
     #and randomly left - right flipping during training.
     input_shape = tf.shape(image)[0:2]
@@ -66,7 +66,7 @@ def rescale_image_and_annotation_by_factor(image, annotation, nin_scale=0.5, max
     annotation = tf.image.resize_images(annotation, scaled_input_shape,
                                         method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
 
-    return image, annotation
+    return image, annotation, image_shape
 
 
 def download_resnet_checkpoint_if_necessary(resnet_checkpoints_path, resnet_model_name):
@@ -93,7 +93,7 @@ def download_resnet_checkpoint_if_necessary(resnet_checkpoints_path, resnet_mode
         print("ResNet checkpoints file successfully found.")
 
 
-def scale_image_with_crop_padding(image, annotation, crop_size):
+def scale_image_with_crop_padding(image, annotation, image_shape, crop_size):
 
     image_croped = tf.image.resize_image_with_crop_or_pad(image,crop_size,crop_size)
 
@@ -108,7 +108,7 @@ def scale_image_with_crop_padding(image, annotation, crop_size):
     annotation_additional_mask_out = tf.to_int32(tf.equal(cropped_padded_annotation, 0)) * (mask_out_number+1)
     cropped_padded_annotation = cropped_padded_annotation + annotation_additional_mask_out - 1
 
-    return image_croped, tf.squeeze(cropped_padded_annotation)
+    return image_croped, tf.squeeze(cropped_padded_annotation), image_shape
 
 def tf_record_parser(record):
     keys_to_features = {
@@ -131,9 +131,9 @@ def tf_record_parser(record):
     annotation = tf.reshape(annotation, (height,width,1), name="annotation_reshape")
     annotation = tf.to_int32(annotation)
 
-    return tf.to_float(image), annotation
+    return tf.to_float(image), annotation, (height, width)
 
-def distort_randomly_image_color(image_tensor, annotation_tensor):
+def distort_randomly_image_color(image_tensor, annotation_tensor, image_shape):
     """Accepts image tensor of (width, height, 3) and returns color distorted image.
     The function performs random brightness, saturation, hue, contrast change as it is performed
     for inception model training in TF-Slim (you can find the link below in comments). All the
@@ -167,4 +167,4 @@ def distort_randomly_image_color(image_tensor, annotation_tensor):
 
     img_float_distorted_original_range = distorted_image * 255
 
-    return img_float_distorted_original_range, annotation_tensor
+    return img_float_distorted_original_range, annotation_tensor, image_shape

test.py

@@ -16,7 +16,7 @@
 parser = argparse.ArgumentParser()
 
 envarg = parser.add_argument_group('Eval params')
-envarg.add_argument("--model_id", type=int, help="Model id name to be loaded.")
+envarg.add_argument("--model_id", default=16645, type=int, help="Model id name to be loaded.")
 input_args = parser.parse_args()
 
 # best: 16645
@@ -68,13 +68,13 @@ class Dotdict(dict):
 class_labels[-1] = 255
 
 LOG_FOLDER = './tboard_logs'
-TEST_DATASET_DIR="./dataset/"
+TEST_DATASET_DIR="./dataset/tfrecords"
 TEST_FILE = 'test.tfrecords'
 
 test_filenames = [os.path.join(TEST_DATASET_DIR,TEST_FILE)]
 test_dataset = tf.data.TFRecordDataset(test_filenames)
 test_dataset = test_dataset.map(tf_record_parser)  # Parse the record into tensors.
-test_dataset = test_dataset.map(scale_image_with_crop_padding)
+test_dataset = test_dataset.map(lambda image, annotation, image_shape: scale_image_with_crop_padding(image, annotation, image_shape, args.crop_size))
 test_dataset = test_dataset.shuffle(buffer_size=100)
 test_dataset = test_dataset.batch(args.batch_size)
 
@@ -88,8 +88,8 @@ class Dotdict(dict):
     logits_batch_tensor=logits_tf,
     class_labels=class_labels)
 
-cross_entropies_tf = tf.nn.softmax_cross_entropy_with_logits(logits=valid_logits_batch_tf,
-                                                             labels=valid_labels_batch_tf)
+cross_entropies_tf = tf.nn.softmax_cross_entropy_with_logits_v2(logits=valid_logits_batch_tf,
+                                                                labels=valid_labels_batch_tf)
 
 cross_entropy_mean_tf = tf.reduce_mean(cross_entropies_tf)
 tf.summary.scalar('cross_entropy', cross_entropy_mean_tf)
@@ -154,12 +154,12 @@ class Dotdict(dict):
                 mean_IoU.append(IoU)
                 mean_freq_weighted_IU.append(freq_weighted_IU)
 
-                #f, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(8, 8))
+                f, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(8, 8))
 
-                #ax1.imshow(input_image.astype(np.uint8))
-                #ax2.imshow(label_image)
-                #ax3.imshow(pred_image)
-                #plt.show()
+                ax1.imshow(input_image.astype(np.uint8))
+                ax2.imshow(label_image)
+                ax3.imshow(pred_image)
+                plt.show()
 
         except tf.errors.OutOfRangeError:
             break

train.py

@@ -68,7 +68,7 @@
 training_dataset = training_dataset.map(tf_record_parser)
 training_dataset = training_dataset.map(rescale_image_and_annotation_by_factor)
 training_dataset = training_dataset.map(distort_randomly_image_color)
-training_dataset = training_dataset.map(lambda image, annotation: scale_image_with_crop_padding(image, annotation, crop_size))
+training_dataset = training_dataset.map(lambda image, annotation, image_shape: scale_image_with_crop_padding(image, annotation, image_shape, crop_size))
 training_dataset = training_dataset.map(random_flip_image_and_annotation)  # Parse the record into tensors.
 training_dataset = training_dataset.repeat()  # number of epochs
 training_dataset = training_dataset.shuffle(buffer_size=500)
@@ -77,7 +77,7 @@
 validation_filenames = [os.path.join(TRAIN_DATASET_DIR,VALIDATION_FILE)]
 validation_dataset = tf.data.TFRecordDataset(validation_filenames)
 validation_dataset = validation_dataset.map(tf_record_parser)  # Parse the record into tensors.
-validation_dataset = validation_dataset.map(lambda image, annotation: scale_image_with_crop_padding(image, annotation, crop_size))
+validation_dataset = validation_dataset.map(lambda image, annotation, image_shape: scale_image_with_crop_padding(image, annotation, image_shape, crop_size))
 validation_dataset = validation_dataset.shuffle(buffer_size=100)
 validation_dataset = validation_dataset.batch(args.batch_size)
 
@@ -92,7 +92,7 @@
 
 iterator = tf.data.Iterator.from_string_handle(
     handle, training_dataset.output_types, training_dataset.output_shapes)
-batch_images_tf, batch_labels_tf = iterator.get_next()
+batch_images_tf, batch_labels_tf, _ = iterator.get_next()
 
 # You can use feedable iterators with a variety of different kinds of iterator
 # (such as one-shot and initializable iterators).