File size: 5,591 Bytes
169a39f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """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 # infer.py calls load_images(..., 25.0)
def load_inverse_depth(self, depth, depth_scale: float, max_depth: float):
# infer.py load_images; a path is the raw camera image (cv2 IMREAD_UNCHANGED, / depth_scale), an array is meters
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 # Remove values beyond max range
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]) # infer.py load_images
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)
# RGBDDepth.image2tensor
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: # infer_image
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"]
|