update pipeline

models/depth_normal_pipeline_clip.py

@@ -26,6 +26,32 @@ from utils.normal_ensemble import ensemble_normals
 from utils.batch_size import find_batch_size
 import cv2
 
+def get_save_width_height(input_image_size, depth_pred):
+  orig_width, orig_height = input_image_size
+  depth_height, depth_width = depth_pred.shape[:2]
+
+  print("Inside pipe!")
+  print("Depth size: {}x{}".format(depth_width, depth_height))
+
+  # Use height
+  ratio = orig_height / depth_height
+  new_width = depth_width * ratio
+  if new_width <= orig_width:
+      save_width = int(new_width)
+      save_height = int(orig_height)
+  else:
+      ratio = orig_width / depth_width
+      new_height = depth_height * ratio
+      if new_height <= orig_height:
+          save_width = int(orig_width)
+          save_height = int(new_height)
+      else:
+          print("WTF IS THIS??")
+          return None, None
+
+  print("New depth size: {}x{}".format(save_width, save_height))
+  return save_width, save_height
+  
 class DepthNormalPipelineOutput(BaseOutput):
     """
     Output class for Marigold monocular depth prediction pipeline.
@@ -179,10 +205,11 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
 
         # Resize back to original resolution
         if match_input_res:
+            save_width, save_height = get_save_width_height(input_size, depth_pred)
             pred_img = Image.fromarray(depth_pred)
-            pred_img = pred_img.resize(input_size)
+            pred_img = pred_img.resize((save_width, save_height))
             depth_pred = np.asarray(pred_img)
-            normal_pred = cv2.resize(chw2hwc(normal_pred), input_size, interpolation = cv2.INTER_NEAREST)
+            normal_pred = cv2.resize(chw2hwc(normal_pred), (save_width, save_height), interpolation = cv2.INTER_NEAREST)
 
         # Clip output range: current size is the original size
         depth_pred = depth_pred.clip(0, 1)