deeper pass

1 files changed, 42 insertions, 5 deletions

diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py
index e388073..ee4f38a 100644
--- a/cli/app/search/search_dense.py
+++ b/cli/app/search/search_dense.py
@@ -216,23 +216,50 @@ 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)
 
+    feat_loss_aa, feat_err_aa = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, 0, 0, height/3, width/3, height, width)
+    feat_loss_ab, feat_err_ab = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, height*1/3, 0, height/3, width/3, height, width)
+    feat_loss_ac, feat_err_ac = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, height*2/3, 0, height/3, width/3, height, width)
+    feat_loss_ad, feat_err_ad = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, 0, width*1/3, height/3, width/3, height, width)
+    feat_loss_ae, feat_err_ae = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, height*1/3, width*1/3, height/3, width/3, height, width)
+    feat_loss_af, feat_err_af = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, height*2/3, width*1/3, height/3, width/3, height, width)
+    feat_loss_ag, feat_err_ag = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, 0, width*2/3, height/3, width/3, height, width)
+    feat_loss_ah, feat_err_ah = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, height*1/3, width*2/3, height/3, width/3, height, width)
+    feat_loss_ai, feat_err_ai = feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, gen_img_ch, target_img_ch, height*2/3, width*2/3, height/3, width/3, height, width)
+
     mse_loss_a = mse_loss_crop(target_img_ch, gen_img_ch, 0, 0, img_w / 2, img_w / 2)
     mse_loss_b = mse_loss_crop(target_img_ch, gen_img_ch, img_w / 2, 0, img_w / 2, img_w / 2)
     mse_loss_c = mse_loss_crop(target_img_ch, gen_img_ch, 0, img_w / 2, img_w / 2, img_w / 2)
     mse_loss_d = mse_loss_crop(target_img_ch, gen_img_ch, img_w / 2, img_w / 2, img_w / 2, img_w / 2)
 
+    mse_loss_aa = mse_loss_crop(target_img_ch, gen_img_ch, 0, 0, img_w/3, img_w/3)
+    mse_loss_ab = mse_loss_crop(target_img_ch, gen_img_ch, img_w*1/3, 0, img_w/3, img_w/3)
+    mse_loss_ac = mse_loss_crop(target_img_ch, gen_img_ch, img_w*2/3, 0, img_w/3, img_w/3)
+    mse_loss_ad = mse_loss_crop(target_img_ch, gen_img_ch, 0, img_w*1/3, img_w/3, img_w/3)
+    mse_loss_ae = mse_loss_crop(target_img_ch, gen_img_ch, img_w*1/3, img_w*1/3, img_w/3, img_w/3)
+    mse_loss_af = mse_loss_crop(target_img_ch, gen_img_ch, img_w*2/3, img_w*1/3, img_w/3, img_w/3)
+    mse_loss_ag = mse_loss_crop(target_img_ch, gen_img_ch, 0, img_w*2/3, img_w/3, img_w/3)
+    mse_loss_ah = mse_loss_crop(target_img_ch, gen_img_ch, img_w*1/3, img_w*2/3, img_w/3, img_w/3)
+    mse_loss_ai = mse_loss_crop(target_img_ch, gen_img_ch, img_w*2/3, img_w*2/3, img_w/3, img_w/3)
+
     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
+
+    feat_loss_quint = feat_loss_aa + feat_loss_ab + feat_loss_ac + feat_loss_ad + feat_loss_ae + feat_loss_af + feat_loss_ag + feat_loss_ah + feat_loss_ai
+    img_feat_err_quint = feat_err_aa + feat_err_ab + feat_err_ac + feat_err_ad + feat_err_ae + feat_err_af + feat_err_ag + feat_err_ah + feat_err_ai
+    mse_loss_quint = mse_loss_aa + mse_loss_ab + mse_loss_ac + mse_loss_ad + mse_loss_ae + mse_loss_af + mse_loss_ag + mse_loss_ah + mse_loss_ai
   else:
     feat_loss = tf.constant(0.0)
     img_feat_err = tf.constant(0.0)
     feat_loss_quad = tf.constant(0.0)
     img_feat_err_quad = tf.constant(0.0)
+    feat_loss_quint = tf.constant(0.0)
+    img_feat_err_quint = tf.constant(0.0)
 
   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_quad + params.lambda_feat * feat_loss_quad
+  inv_loss = (params.lambda_mse * mse_loss + params.lambda_feat * feat_loss) * 9
+  inv_loss_quad = (params.lambda_mse * mse_loss_quad + params.lambda_feat * feat_loss_quad) * 9/4
+  inv_loss_quint = params.lambda_mse * mse_loss_quint + params.lambda_feat * feat_loss_quint
 
   # --------------------------
   # Optimizer.
@@ -363,10 +390,12 @@ def find_dense_embedding_for_images(params, opt_tag="inverse_" + timestamp(), op
     print("Beginning dense iteration...")
     for _ in range(params.inv_it):
 
-      if it < params.inv_it * 0.75:
+      if it < params.inv_it * 0.5:
         _inv_loss, _mse_loss, _feat_loss, _lrate, _ = sess.run([inv_loss, mse_loss, feat_loss, lrate, inv_train_op])
-      else:
+      elif it < params.inv_it * 0.75:
         _inv_loss, _mse_loss, _feat_loss, _lrate, _ = sess.run([inv_loss_quad, mse_loss, feat_loss_quad, lrate, inv_train_op_quad])
+      else:
+        _inv_loss, _mse_loss, _feat_loss, _lrate, _ = sess.run([inv_loss_quint, mse_loss, feat_loss_quint, lrate, inv_train_op_quint])
 
       if params.clipping or params.stochastic_clipping:
         sess.run(clip_latent)
@@ -459,14 +488,22 @@ def mse_loss_crop(img_a, img_b, 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):
+def feature_loss(feature_extractor, opt_feature_layers, BATCH_SIZE, img_a, img_b, y, x, height, width, resize_height=None, resize_width=None):
+  height = int(height)
+  width = int(width)
   if y is not None:
+    x = int(x)
+    y = int(y)
     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)
   else:
     img_a = tf.image.resize_images(img_a, [height, width])
     img_b = tf.image.resize_images(img_b, [height, width])
 
+  if resize_height is not None:
+    img_a = tf.image.resize_images(img_a, [resize_height, resize_width])
+    img_b = tf.image.resize_images(img_b, [resize_height, resize_width])
+
   gen_feat_ex = feature_extractor(dict(images=img_a), as_dict=True, signature='image_feature_vector')
   target_feat_ex = feature_extractor(dict(images=img_b), as_dict=True, signature='image_feature_vector')