Create app.py

app.py

@@ -0,0 +1,46 @@
+import torch
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torchvision.transforms as T
+from model import HybridDepthModel
+
+# Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = HybridDepthModel().to(device)
+model.load_state_dict(torch.load("depth_model_all.pth", map_location=device))
+model.eval()
+
+# Preprocessing
+transform = T.Compose([
+    T.Resize((32, 32)),
+    T.ToTensor(),
+])
+
+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]
+
+# Gradio UI
+examples = [["example.png"]]
+
+demo = gr.Interface(
+    fn=predict_depth,
+    inputs=gr.Image(type="pil", label="Input RGB Image"),
+    outputs=[
+        gr.Image(type="pil", label="Original Image"),
+        gr.Image(type="pil", label="Predicted Depth Map"),
+    ],
+    title="🔭 DepthStar: Light-weight Depth Estimation",
+    description="Upload an RGB image and get the depth map predicted by our tiny DepthStar model.",
+    examples=examples,
+    allow_flagging="never",
+    theme="huggingface",
+)
+
+if __name__ == "__main__":
+    demo.launch()