File size: 3,269 Bytes
397f342 | 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 | """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"]
|