| """Depth Anything V2 preprocessing: the official ``DepthAnythingV2.image2tensor`` transform (Resize lower_bound to |
| 518 keeping the aspect ratio, sides multiples of 14, cv2 bicubic; ImageNet normalisation; vendored in |
| ``dav2_dpt.py``) behind the transformers image-processor interface. A path is read with cv2 as the official scripts |
| do; a PIL image / RGB array is used as RGB. ``post_process_depth_estimation`` resizes to the given sizes as |
| ``infer_image`` does (bilinear, align_corners=True -- not transformers' bicubic).""" |
|
|
| 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 .dav2_dpt import NormalizeImage, PrepareForNet, Resize |
|
|
|
|
| class DepthAnythingV2ImageProcessor(BaseImageProcessor): |
| model_input_names = ["pixel_values"] |
|
|
| def __init__(self, input_size: int = 518, **kwargs): |
| super().__init__(**kwargs) |
| self.input_size = input_size |
|
|
| def preprocess(self, images, input_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature: |
| if not isinstance(images, (list, tuple)): |
| images = [images] |
| size = input_size or self.input_size |
| transform = Compose([ |
| Resize(width=size, height=size, resize_target=False, 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(), |
| ]) |
| tensors = [] |
| for image in images: |
| if isinstance(image, str): |
| rgb = cv2.cvtColor(cv2.imread(image), cv2.COLOR_BGR2RGB) / 255.0 |
| else: |
| rgb = np.asarray(image.convert("RGB") if isinstance(image, Image.Image) else image) / 255.0 |
| tensors.append(torch.from_numpy(transform({"image": rgb})["image"])) |
| if len({t.shape for t in tensors}) > 1: |
| raise ValueError("images resize to different shapes; process them one at a time") |
| return BatchFeature({"pixel_values": torch.stack(tensors)}, tensor_type=return_tensors) |
|
|
| def post_process_depth_estimation(self, outputs, target_sizes=None) -> list: |
| """``DepthEstimatorOutput`` -> ``[{"predicted_depth": (h, w)}]`` 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="bilinear", align_corners=True)[0, 0] |
| results.append({"predicted_depth": depth}) |
| return results |
|
|
|
|
| __all__ = ["DepthAnythingV2ImageProcessor"] |
|
|
