bugs

app.py

@@ -34,7 +34,7 @@ else:
 
 
 def estimate_depth(image: Image) -> Image:
-    image= image.resize((384,384))
+    image= image.resize((384,384)).convert("RGB")
     image = np.array(image)
     input_batch = transform(image).to(device)
 
@@ -51,6 +51,11 @@ def estimate_depth(image: Image) -> Image:
     output = prediction.cpu().numpy()
 
     output= 255 * output/np.max(output)
+
+    # If 3 channels convert to 1
+    if (len(output.shape) == 3):
+        if (output.shape[2] == 3):
+            output = output[:, :, 0]
     
     return Image.fromarray(output.astype("uint8"))