RGBD-SOD/dptdepth

1from typing import Dict, Optional, Tuple2 3import numpy as np4import torch.nn.functional as F5import torchvision.transforms as transforms6import torchvision.transforms.functional as TF7from PIL.Image import Image8from torch import Tensor9from transformers.image_processing_utils import BaseImageProcessor10 11INPUT_IMAGE_SIZE = (352, 352)12 13transform = transforms.Compose(14    [15        transforms.Resize(16            INPUT_IMAGE_SIZE,17            interpolation=TF.InterpolationMode.BICUBIC,18        ),19        transforms.ToTensor(),20        transforms.Normalize(21            (0.5, 0.5, 0.5),22            (0.5, 0.5, 0.5),23        ),24    ]25)26 27 28class DPTDepthImageProcessor(BaseImageProcessor):29    model_input_names = ["dptdepth_preprocessor"]30 31    def __init__(self, testsize: Optional[int] = 352, **kwargs) -> None:32        super().__init__(**kwargs)33        self.testsize = testsize34 35    def preprocess(36        self, inputs: Dict[str, Image], **kwargs  # {'rgb': ... }37    ) -> Dict[str, Tensor]:38        rgb: Tensor = transform(inputs["rgb"])39        return dict(rgb=rgb.unsqueeze(0))40 41    def postprocess(42        self, logits: Tensor, size: Tuple[int, int], **kwargs43    ) -> np.ndarray:44        logits: Tensor = F.upsample(45            logits, size=size, mode="bilinear", align_corners=False46        )47        res: np.ndarray = logits.squeeze().data.cpu().numpy()48        # res = (res - res.min()) / (res.max() - res.min() + 1e-8)49        return res50