fix bugs in evaluation

ychfan · Mar 30, 2017 · ddfc5de · ddfc5de
1 parent 22d59af
commit ddfc5de
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/predict.py b/predict.py
@@ -28,10 +28,13 @@
     lr_image = tf.image.decode_image(image_file, channels=3)
     lr_image = tf.image.convert_image_dtype(lr_image, tf.float32)
     lr_image = tf.expand_dims(lr_image, 0)
-    hr_image_shape = tf.shape(lr_image)[1:3] * FLAGS.scale
+    lr_image_shape = tf.shape(lr_image)[1:3]
+    hr_image_shape = lr_image_shape * FLAGS.scale
     if (data.resize_func is not None):
         lr_image = data.resize_func(lr_image, hr_image_shape)
-    lr_image = tf.reshape(lr_image, [1, hr_image_shape[0], hr_image_shape[1], 3])
+        lr_image = tf.reshape(lr_image, [1, hr_image_shape[0], hr_image_shape[1], 3])
+    else:
+        lr_image = tf.reshape(lr_image, [1, lr_image_shape[0], lr_image_shape[1], 3])
     lr_image = util.pad_boundary(lr_image)
     hr_image = model.build_model(lr_image, FLAGS.scale, False)
     hr_image = util.crop_center(hr_image, hr_image_shape)

diff --git a/validate.py b/validate.py
@@ -31,10 +31,13 @@
     lr_image = tf.image.convert_image_dtype(lr_image, tf.float32)
     hr_image = tf.expand_dims(hr_image, 0)
     lr_image = tf.expand_dims(lr_image, 0)
+    lr_image_shape = tf.shape(lr_image)[1:3]
     hr_image_shape = tf.shape(hr_image)[1:3]
     if (data.resize_func is not None):
         lr_image = data.resize_func(lr_image, hr_image_shape)
-    lr_image = tf.reshape(lr_image, [1, hr_image_shape[0], hr_image_shape[1], 3])
+        lr_image = tf.reshape(lr_image, [1, hr_image_shape[0], hr_image_shape[1], 3])
+    else:
+        lr_image = tf.reshape(lr_image, [1, lr_image_shape[0], lr_image_shape[1], 3])
     lr_image = util.pad_boundary(lr_image)
     lr_image = model.build_model(lr_image, FLAGS.scale, False)
     lr_image = util.crop_center(lr_image, hr_image_shape)