| """Video Depth Anything preprocessing: the per-frame transform of the official ``infer_video_depth`` (Resize lower_bound |
| to 518 keeping the aspect ratio, sides multiples of 14, cv2 bicubic; ImageNet normalisation; vendored in ``vda_dpt.py``), |
| including its rule for wide videos (aspect ratio > 1.78: smaller input size), behind the transformers image-processor |
| interface. Frames are RGB uint8 arrays (H, W, 3), as the official ``read_video_frames`` returns them. |
| ``post_process_depth_estimation`` resizes to the given sizes as ``infer_video_depth`` does (float32, bilinear, |
| align_corners=True).""" |
|
|
| 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 .vda_dpt import NormalizeImage, PrepareForNet, Resize |
|
|
|
|
| class VideoDepthAnythingImageProcessor(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, videos, input_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature: |
| """videos: one video or a list of videos of the same size; a video is an array (T, H, W, 3) or a list of |
| frames (RGB arrays or PIL images). pixel_values: (B, T, 3, H', W').""" |
| is_frame = lambda x: isinstance(x, Image.Image) or isinstance(x, np.ndarray) and x.ndim == 3 |
| if isinstance(videos, np.ndarray) and videos.ndim == 4 or is_frame(videos[0]): |
| videos = [videos] |
| size = input_size or self.input_size |
| batch = [] |
| for video in videos: |
| frames = [np.asarray(f.convert("RGB")) if isinstance(f, Image.Image) else f for f in video] |
| frame_height, frame_width = frames[0].shape[:2] |
| |
| input_size_video = size |
| ratio = max(frame_height, frame_width) / min(frame_height, frame_width) |
| if ratio > 1.78: |
| input_size_video = int(input_size_video * 1.777 / ratio) |
| input_size_video = round(input_size_video / 14) * 14 |
| transform = Compose([ |
| Resize(width=input_size_video, height=input_size_video, 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(), |
| ]) |
| batch.append(torch.stack([torch.from_numpy(transform({"image": f.astype(np.float32) / 255.0})["image"]) for f in frames])) |
| if len({t.shape for t in batch}) > 1: |
| raise ValueError("videos resize to different shapes; process them one at a time") |
| return BatchFeature({"pixel_values": torch.stack(batch)}, tensor_type=return_tensors) |
|
|
| def post_process_depth_estimation(self, outputs, target_sizes=None) -> list: |
| """``DepthEstimatorOutput`` (predicted_depth (B, T, H, W)) -> ``[{"predicted_depth": (T, h, w)}]`` per video; |
| target_sizes: (height, width) per video (e.g. the frame size), 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.float().unsqueeze(1), size=(int(size[0]), int(size[1])), mode="bilinear", align_corners=True)[:, 0] |
| results.append({"predicted_depth": depth}) |
| return results |
|
|
|
|
| __all__ = ["VideoDepthAnythingImageProcessor"] |
|
|
