| """Camera Depth Model preprocessing: cdm/infer.py ``load_images`` (raw depth / depth_scale, values above max_depth set |
| to 0, inverse depth where valid) and ``RGBDDepth.image2tensor`` (Resize lower_bound to 518 keeping the aspect ratio, |
| sides multiples of 14, cv2 bicubic for the image and nearest for the depth; ImageNet normalisation; vendored in |
| ``cdm_dpt.py``) behind the transformers image-processor interface. |
| |
| As in infer.py and the README API, the image is given as RGB and ``image2tensor`` applies ``cv2.COLOR_BGR2RGB`` to it, |
| so the network sees the channels swapped; this is kept unchanged. ``post_process_depth_estimation`` resizes to the |
| given sizes as ``infer_image`` does (nearest). The model predicts inverse depth: meters = 1 / prediction (infer.py).""" |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image |
| from torchvision.transforms import Compose |
| from transformers.feature_extraction_utils import BatchFeature |
| from transformers.image_processing_utils import BaseImageProcessor |
|
|
| from .cdm_dpt import NormalizeImage, PrepareForNet, Resize |
|
|
|
|
| class CameraDepthModelImageProcessor(BaseImageProcessor): |
| model_input_names = ["pixel_values", "prompt_depth"] |
|
|
| def __init__(self, input_size: int = 518, depth_scale: float = 1000.0, max_depth: float = 25.0, **kwargs): |
| super().__init__(**kwargs) |
| self.input_size = input_size |
| self.depth_scale = depth_scale |
| self.max_depth = max_depth |
|
|
| def load_inverse_depth(self, depth, depth_scale: float, max_depth: float): |
| |
| if isinstance(depth, str): |
| depth_low_res = cv2.imread(depth, cv2.IMREAD_UNCHANGED) |
| if depth_low_res is None: |
| raise ValueError(f"Could not load depth image from {depth}") |
| depth_low_res = np.asarray(depth_low_res).astype(np.float32) / depth_scale |
| else: |
| depth_low_res = np.array(depth, dtype=np.float32) |
| depth_low_res[depth_low_res > max_depth] = 0.0 |
| simi_depth_low_res = np.zeros_like(depth_low_res) |
| simi_depth_low_res[depth_low_res > 0] = 1 / depth_low_res[depth_low_res > 0] |
| return simi_depth_low_res |
|
|
| def preprocess(self, images, depth=None, input_size: int | None = None, depth_scale: float | None = None, |
| max_depth: float | None = None, return_tensors=None, **kwargs) -> BatchFeature: |
| """images: one image or a list (paths, PIL images or RGB uint8 arrays); depth: one raw depth per image (paths |
| of the camera's depth png, or arrays in meters). pixel_values (B, 3, H, W); prompt_depth (B, 1, H, W): the |
| inverse depth resized with the image.""" |
| if not isinstance(images, (list, tuple)): |
| images = [images] |
| if not isinstance(depth, (list, tuple)): |
| depth = [depth] |
| if len(depth) != len(images): |
| raise ValueError("give one depth per image") |
| size = input_size or self.input_size |
| transform = Compose([ |
| Resize(width=size, height=size, resize_target=True, keep_aspect_ratio=True, ensure_multiple_of=14, |
| resize_method="lower_bound", image_interpolation_method=cv2.INTER_CUBIC), |
| NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
| PrepareForNet(), |
| ]) |
| pixel_values, prompt_depth = [], [] |
| for image, raw_depth in zip(images, depth): |
| if isinstance(image, str): |
| raw_image = np.asarray(cv2.imread(image)[:, :, ::-1]) |
| else: |
| raw_image = np.asarray(image.convert("RGB") if isinstance(image, Image.Image) else image) |
| simi_depth = self.load_inverse_depth(raw_depth, depth_scale or self.depth_scale, max_depth or self.max_depth) |
| |
| prepared = transform({"image": cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0, "depth": simi_depth}) |
| pixel_values.append(torch.from_numpy(prepared["image"])) |
| prompt_depth.append(torch.from_numpy(prepared["depth"]).unsqueeze(0)) |
| if len({t.shape for t in pixel_values}) > 1: |
| raise ValueError("images resize to different shapes; process them one at a time") |
| return BatchFeature({"pixel_values": torch.stack(pixel_values), "prompt_depth": torch.stack(prompt_depth)}, |
| tensor_type=return_tensors) |
|
|
| def post_process_depth_estimation(self, outputs, target_sizes=None) -> list: |
| """``DepthEstimatorOutput`` -> ``[{"predicted_depth": (h, w)}]`` (inverse depth) per image; target_sizes: |
| (height, width) per image (e.g. the original sizes), None keeps the model resolution.""" |
| predicted_depth = outputs.predicted_depth |
| if target_sizes is not None and len(predicted_depth) != len(target_sizes): |
| raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the predicted depth") |
| results = [] |
| for depth, size in zip(predicted_depth, target_sizes if target_sizes is not None else [None] * len(predicted_depth)): |
| if size is not None: |
| depth = F.interpolate(depth[None, None], (int(size[0]), int(size[1])), mode="nearest")[0, 0] |
| results.append({"predicted_depth": depth}) |
| return results |
|
|
|
|
| __all__ = ["CameraDepthModelImageProcessor"] |
|
|
