Views
No views yet
1from typing import Dict
2
3import numpy as np
4from datasets import load_dataset
5from matplotlib import cm
6from PIL import Image
7from torch import Tensor
8from transformers import AutoImageProcessor, AutoModel
9
10model = AutoModel.from_pretrained("RGBD-SOD/bbsnet", trust_remote_code=True)
11image_processor = AutoImageProcessor.from_pretrained(
12 "RGBD-SOD/bbsnet", trust_remote_code=True
13)
14dataset = load_dataset("RGBD-SOD/test", "v1", split="train", cache_dir="data")
15
16index = 0
17
18"""
19Get a specific sample from the dataset
20
21sample = {
22 'depth': <PIL.PngImagePlugin.PngImageFile image mode=L size=640x360>,
23 'rgb': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=640x360>,
24 'gt': <PIL.PngImagePlugin.PngImageFile image mode=L size=640x360>,
25 'name': 'COME_Train_5'
26}
27"""
28sample = dataset[index]
29
30depth: Image.Image = sample["depth"]
31rgb: Image.Image = sample["rgb"]
32gt: Image.Image = sample["gt"]
33name: str = sample["name"]
34
35
36"""
371. Preprocessing step
38
39preprocessed_sample = {
40 'rgb': tensor([[[[-0.8507, ....0365]]]]),
41 'gt': tensor([[[[0., 0., 0...., 0.]]]]),
42 'depth': tensor([[[[0.9529, 0....3490]]]])
43}
44"""
45preprocessed_sample: Dict[str, Tensor] = image_processor.preprocess(sample)
46
47"""
482. Prediction step
49
50output = {
51 'logits': tensor([[[[-5.1966, ...ackward0>)
52}
53"""
54output: Dict[str, Tensor] = model(
55 preprocessed_sample["rgb"], preprocessed_sample["depth"]
56)
57
58"""
593. Postprocessing step
60"""
61postprocessed_sample: np.ndarray = image_processor.postprocess(
62 output["logits"], [sample["gt"].size[1], sample["gt"].size[0]]
63)
64prediction = Image.fromarray(np.uint8(cm.gist_earth(postprocessed_sample) * 255))
65
66"""
67Show the predicted salient map and the corresponding ground-truth(GT)
68"""
69prediction.show()
70gt.show()
71@inproceedings{fan2020bbs,
title={BBS-Net: RGB-D salient object detection with a bifurcated backbone strategy network},
author={Fan, Deng-Ping and Zhai, Yingjie and Borji, Ali and Yang, Jufeng and Shao, Ling},
booktitle={Computer Vision--ECCV 2020: 16th European Conference, Glasgow, UK, August 23--28, 2020, Proceedings, Part XII},
pages={275--292},
year={2020},
organization={Springer}
}