Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 from PIL import Image
 import torchvision.transforms as T
 from model import DepthSTAR
+import matplotlib.cm as cm
 
 # Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -24,17 +25,24 @@ def predict_depth(image):
     img = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
         pred = model(img)[0, 0].cpu().numpy()
-    
+
+    # Normalize prediction
     pred_normalized = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
-    pred_image = Image.fromarray((pred_normalized * 255).astype(np.uint8))
 
-    # Resize both for better display (e.g., 512x512)
+    # Apply inferno colormap
+    cmap = cm.get_cmap('inferno')
+    pred_colored = cmap(pred_normalized)[:, :, :3]  # Drop alpha channel
+    pred_colored = (pred_colored * 255).astype(np.uint8)
+    pred_pil = Image.fromarray(pred_colored)
+
+    # Upscale for display
     upscale_size = (256, 256)
     image_resized = image.resize(upscale_size, resample=Image.NEAREST)
-    pred_resized = pred_image.resize(upscale_size, resample=Image.NEAREST)
+    pred_resized = pred_pil.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"],]
 examples = [["example.png"]]