Init

.gitignore

@@ -0,0 +1,8 @@
+env
+__pycache__
+data
+samples/*
+!samples/.gitkeep
+model_cache
+image_processor_cache
+flagged

app.py

@@ -0,0 +1,25 @@
+import gradio as gr
+
+from inference import inference
+from prepare_samples import prepare_samples
+
+TITLE = "DPT Depth"
+DESCRIPTION = "Gradio demo for DPT-Depth"
+examples = prepare_samples()
+
+demo = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.inputs.Image(label="RGB", type="pil"),
+    ],
+    outputs=[
+        gr.outputs.Image(label="Prediction", type="pil"),
+    ],
+    title=TITLE,
+    examples=examples,
+    description=DESCRIPTION,
+)
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")  # server_port=8541

inference.py

@@ -0,0 +1,35 @@
+from transformers import AutoImageProcessor, AutoModel
+from typing import Dict
+
+import numpy as np
+from matplotlib import cm
+from PIL import Image
+from torch import Tensor
+
+model = AutoModel.from_pretrained(
+    "RGBD-SOD/dptdepth", trust_remote_code=True, cache_dir="model_cache"
+)
+image_processor = AutoImageProcessor.from_pretrained(
+    "RGBD-SOD/dptdepth", trust_remote_code=True, cache_dir="image_processor_cache"
+)
+
+
+def inference(rgb: Image.Image) -> Image.Image:
+    rgb = rgb.convert(mode="RGB")
+
+    preprocessed_sample: Dict[str, Tensor] = image_processor.preprocess(
+        {
+            "rgb": rgb,
+        }
+    )
+
+    output: Dict[str, Tensor] = model(preprocessed_sample["rgb"])
+    postprocessed_sample: np.ndarray = image_processor.postprocess(
+        output["logits"], [rgb.size[1], rgb.size[0]]
+    )
+    prediction = Image.fromarray(np.uint8(cm.gist_earth(postprocessed_sample) * 255))
+    return prediction
+
+
+if __name__ == "__main__":
+    pass

prepare_samples.py

@@ -0,0 +1,31 @@
+import os
+import shutil
+from typing import List, Tuple
+
+from PIL import Image
+from datasets import load_dataset
+
+
+dataset = load_dataset("RGBD-SOD/test", "v1", split="train", cache_dir="data")
+SAMPLES_DIR = "samples"
+
+
+def prepare_samples():
+    samples: List[Tuple[str, str, str]] = []
+    for sample in dataset:
+        rgb: Image.Image = sample["rgb"]
+        depth: Image.Image = sample["depth"]
+        gt: Image.Image = sample["gt"]
+        name: str = sample["name"]
+        dir_path = os.path.join(SAMPLES_DIR, name)
+        shutil.rmtree(dir_path, ignore_errors=True)
+        os.makedirs(dir_path, exist_ok=True)
+        rgb_path = os.path.join(dir_path, f"rgb.jpg")
+        rgb.save(rgb_path)
+        depth_path = os.path.join(dir_path, f"depth.jpg")
+        depth.save(depth_path)
+        gt_path = os.path.join(dir_path, f"gt.png")
+        gt.save(gt_path)
+
+        samples.append([rgb_path, depth_path, gt_path])
+    return samples

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+torch
+opencv-python
+transformers[torch]
+torchvision
+datasets
+matplotlib
+timm