Update app.py

app.py

@@ -91,7 +91,7 @@ def predict_images(dict, depth, prompt="", negative_prompt="", guidance_scale=7.
     if negative_prompt == "":
         negative_prompt = None
 
-    og_size = dict["image"].shape[:2]
+    og_size = (dict["image"].shape[1], dict["image"].shape[0])
 
     init_image = cv2.resize(dict["image"], (512, 512))
 
@@ -201,6 +201,7 @@ def predict_images_3d(dict, depth, prompt="", negative_prompt="", guidance_scale
     if negative_prompt == "":
         negative_prompt = None
 
+    og_size = (dict["image"].shape[1], dict["image"].shape[0])
     init_image = cv2.resize(dict["image"], (512, 512))
 
     mask = Image.fromarray(cv2.resize(dict["mask"], (512, 512))[:,:,0])
@@ -221,8 +222,12 @@ def predict_images_3d(dict, depth, prompt="", negative_prompt="", guidance_scale
     
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, depth_image=depth_image, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
     
+    # resize to original size
+    depth_image = depth_image.resize(og_size)
+    output_depth = output.depth[0].resize(og_size)
+    
     depth_in = denormalize(np.array(depth_image), min, max)
-    depth_out = denormalize(np.array(output.depth[0]), min, max)
+    depth_out = denormalize(np.array(output_depth), min, max)
     
     output_image = output.rgb[0]