RGBD-SOD/bbsnet

1from transformers import AutoImageProcessor, AutoModel2from typing import Dict3 4import numpy as np5from matplotlib import cm6from PIL import Image7from torch import Tensor8 9model = AutoModel.from_pretrained(10    "RGBD-SOD/bbsnet", trust_remote_code=True, cache_dir="model_cache"11)12image_processor = AutoImageProcessor.from_pretrained(13    "RGBD-SOD/bbsnet", trust_remote_code=True, cache_dir="image_processor_cache"14)15 16 17def inference(rgb: Image.Image, depth: Image.Image) -> Image.Image:18    rgb = rgb.convert(mode="RGB")19    depth = depth.convert(mode="L")20 21    preprocessed_sample: Dict[str, Tensor] = image_processor.preprocess(22        {23            "rgb": rgb,24            "depth": depth,25        }26    )27 28    output: Dict[str, Tensor] = model(29        preprocessed_sample["rgb"], preprocessed_sample["depth"]30    )31    postprocessed_sample: np.ndarray = image_processor.postprocess(32        output["logits"], [rgb.size[1], rgb.size[0]]33    )34    prediction = Image.fromarray(np.uint8(cm.gist_earth(postprocessed_sample) * 255))35    return prediction36 37 38if __name__ == "__main__":39    pass40