Update app.py

app.py

@@ -24,9 +24,16 @@ def predict_depth(image):
     img = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
         pred = model(img)[0, 0].cpu().numpy()
+    
     pred_normalized = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
-    pred_image = Image.fromarray((pred_normalized * 255).astype(np.uint8))
-    return [image, pred_image]
+    pred_image = Image.fromarray((pred_normalized * 255).astype(np.uint8)).convert("L")
+
+    # Resize both for better display (e.g., 512x512)
+    upscale_size = (512, 512)
+    image_resized = image.resize(upscale_size, resample=Image.NEAREST)
+    pred_resized = pred_image.resize(upscale_size, resample=Image.NEAREST)
+
+    return [image_resized, pred_resized]
 
 # Gradio UI
 examples = [["img_000.png"],["img_001.png"],["img_002.png"],["img_003.png"],["img_004.png"],["img_005.png"],]
@@ -42,7 +49,7 @@ demo = gr.Interface(
     description="Upload an RGB image and get the depth map predicted by our tiny DepthStar model.",
     examples=examples,
     allow_flagging="never",
-    theme="huggingface",
+    theme="darkdefault",
 )
 
 if __name__ == "__main__":