Update

.gitignore

@@ -2,4 +2,5 @@ env
 __pycache__
 data
 samples/*
-!samples/.gitkeep
+!samples/.gitkeep
+model_cache

app.py

@@ -1,22 +1,7 @@
-from typing import Tuple
 import gradio as gr
-import os
-from PIL import Image
-from datasets import load_dataset
-from prepare_samples import prepare_samples
-
-DIR_PATH = os.path.dirname(__file__)
-
-
-def inference(rgb: Image.Image, depth: Image.Image) -> Tuple[Image.Image]:
-    return rgb
 
-
-dataset = load_dataset("RGBD-SOD/test", "v1", split="train", cache_dir="data")
-
-# with gr.Blocks() as demo:
-#     with gr.Row(elem_id="center"):
-#         gr.Markdown("# BBS-Net Demo")
+from inference import inference
+from prepare_samples import prepare_samples
 
 TITLE = "BBS-Net Demo"
 DESCRIPTION = "Gradio demo for BBS-Net: RGB-D salient object detection with a bifurcated backbone strategy network."

inference.py

@@ -0,0 +1,36 @@
+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/bbsnet", trust_remote_code=True, cache_dir="model_cache"
+)
+image_processor = AutoImageProcessor.from_pretrained(
+    "RGBD-SOD/bbsnet", trust_remote_code=True, cache_dir="image_processor_cache"
+)
+
+
+def inference(rgb: Image.Image, depth: Image.Image) -> Image.Image:
+    preprocessed_sample: Dict[str, Tensor] = image_processor.preprocess(
+        {
+            "rgb": rgb,
+            "depth": depth,
+        }
+    )
+
+    output: Dict[str, Tensor] = model(
+        preprocessed_sample["rgb"], preprocessed_sample["depth"]
+    )
+    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

@@ -1,8 +1,10 @@
-from typing import List, Tuple
-from datasets import load_dataset
-from PIL import Image
 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"

requirements.txt

@@ -1,2 +1,7 @@
 gradio
+torch
+opencv-python
+transformers[torch]
+torchvision
 datasets
+matplotlib