adding mse loss terms

1 files changed, 12 insertions, 2 deletions

diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py
index 75bba19..ca91943 100644
--- a/cli/app/search/search_dense.py
+++ b/cli/app/search/search_dense.py
@@ -216,9 +216,14 @@ def find_dense_embedding_for_images(params, opt_tag="inverse_" + timestamp(), op
     feat_loss_d, feat_err_d = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, img_w - width, img_w - width, height, width)
     feat_loss_e, feat_err_e = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, int((img_w - width) / 2), int((img_w - width) / 2), height, width)
 
+    mse_loss_a = mse_loss_crop(target_img, gen_img, 0, 0, height / 2, width / 2)
+    mse_loss_b = mse_loss_crop(target_img, gen_img, height / 2, 0, height / 2, width / 2)
+    mse_loss_c = mse_loss_crop(target_img, gen_img, 0, width / 2, height / 2, width / 2)
+    mse_loss_d = mse_loss_crop(target_img, gen_img, height / 2, width / 2, height / 2, width / 2)
+
     feat_loss_quad = feat_loss_a + feat_loss_b + feat_loss_c + feat_loss_d + feat_loss_e
     img_feat_err_quad = feat_err_a + feat_err_b + feat_err_c + feat_err_d + feat_err_e
-
+    mse_loss_quad = mse_loss_a + mse_loss_b + mse_loss_c + mse_loss_d
   else:
     feat_loss = tf.constant(0.0)
     img_feat_err = tf.constant(0.0)
@@ -227,7 +232,7 @@ def find_dense_embedding_for_images(params, opt_tag="inverse_" + timestamp(), op
 
   img_rec_err = params.lambda_mse * img_mse_err + params.lambda_feat * img_feat_err
   inv_loss = params.lambda_mse * mse_loss + params.lambda_feat * feat_loss
-  inv_loss_quad = params.lambda_mse * mse_loss + params.lambda_feat * feat_loss_quad
+  inv_loss_quad = params.lambda_mse * mse_loss_quad + params.lambda_feat * feat_loss_quad
 
   # --------------------------
   # Optimizer.
@@ -445,6 +450,11 @@ def find_dense_embedding_for_images(params, opt_tag="inverse_" + timestamp(), op
   out_file.close()
   sess.close()
 
+def mse_loss_crop(img_a, img_b, y, x, height, width):
+  img_a = tf.image.crop_to_bounding_box(img_a, y, x, height, width)
+  img_b = tf.image.crop_to_bounding_box(img_b, y, x, height, width)
+  return tf.reduce_mean(tf.square((img_a - img_b) / 2.0))
+
 def feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, img_a, img_b, y, x, height, width):
   if y is not None:
     img_a = tf.image.crop_to_bounding_box(img_a, y, x, height, width)