Add Replicate demo and API

README.md

@@ -1,5 +1,6 @@
 # **ZoeDepth: Combining relative and metric depth** (Official implementation)  <!-- omit in toc -->
 [![Open In Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/isl-org/ZoeDepth)
+[![Replicate](https://replicate.com/cjwbw/zoedepth/badge)](https://replicate.com/cjwbw/zoedepth) 
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT) ![PyTorch](https://img.shields.io/badge/PyTorch_v1.10.1-EE4C2C?&logo=pytorch&logoColor=white) 
 [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/zoedepth-zero-shot-transfer-by-combining/monocular-depth-estimation-on-nyu-depth-v2)](https://paperswithcode.com/sota/monocular-depth-estimation-on-nyu-depth-v2?p=zoedepth-zero-shot-transfer-by-combining)

cog.yaml

@@ -0,0 +1,19 @@
+build:
+  gpu: true
+  python_version: "3.10"
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+  python_packages:
+    - "timm==0.6.12"
+    - "torch==1.12.1"
+    - "torchvision==0.13.1"
+    - "wget==3.2"
+    - "opencv-python==4.7.0.68"
+    - "matplotlib==3.6.3"
+    - "scipy==1.9.3"
+
+  run:
+   - mkdir -p /root/.cache/torch/hub/checkpoints/ && wget --output-document "/root/.cache/torch/hub/checkpoints/ZoeD_M12_N.pt" "https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_N.pt"  && wget --output-document "/root/.cache/torch/hub/checkpoints/ZoeD_M12_K.pt" "https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_K.pt" && wget --output-document "/root/.cache/torch/hub/checkpoints/ZoeD_M12_NK.pt" "https://github.com/isl-org/ZoeDepth/releases/download/v1.0/ZoeD_M12_NK.pt"
+
+predict: "predict.py:Predictor"

predict.py

@@ -0,0 +1,38 @@
+import torch
+from PIL import Image
+from cog import BasePredictor, Input, Path
+
+from zoedepth.utils.misc import get_image_from_url, colorize
+
+
+class Predictor(BasePredictor):
+    def setup(self):
+        """Load the model into memory to make running multiple predictions efficient"""
+        self.models = {
+            k: torch.hub.load(".", k, source="local", pretrained=True).to("cuda")
+            for k in ["ZoeD_N", "ZoeD_K", "ZoeD_NK"]
+        }
+
+    def predict(
+        self,
+        model_type: str = Input(
+            default="ZoeD_N",
+            choices=["ZoeD_N", "ZoeD_K", "ZoeD_NK"],
+        ),
+        image: Path = Input(description="Input image"),
+    ) -> Path:
+        """Run a single prediction on the model"""
+        zoe = self.models[model_type]
+        image = Image.open(image).convert("RGB")  # load
+        # depth_numpy = zoe.infer_pil(image)  # as numpy
+
+        depth_tensor = zoe.infer_pil(image, output_type="tensor")  # as torch tensor
+
+        # Colorize output
+        colored = colorize(depth_tensor)
+
+        # save colored output
+        output_path = "/tmp/out.png"
+        Image.fromarray(colored).save(output_path)
+
+        return Path(output_path)