Update

.gitignore

@@ -0,0 +1,3 @@
+env
+__pycache__
+data

app.py

@@ -0,0 +1,41 @@
+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")
+
+TITLE = "BBS-Net Demo"
+DESCRIPTION = "Gradio demo for BBS-Net: RGB-D salient object detection with a bifurcated backbone strategy network."
+examples = prepare_samples()
+
+demo = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.inputs.Image(label="RGB", type="pil"),
+        gr.inputs.Image(label="Depth", type="pil"),
+    ],
+    outputs=[
+        gr.outputs.Image(label="Prediction", type="pil"),
+    ],
+    title=TITLE,
+    examples=examples,
+    description=DESCRIPTION,
+)
+
+
+if __name__ == "__main__":
+    demo.launch(enable_queue=True, cache_examples=False)

prepare_samples.py

@@ -0,0 +1,29 @@
+from typing import List, Tuple
+from datasets import load_dataset
+from PIL import Image
+import os
+import shutil
+
+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,2 @@
+gradio
+datasets