| """Prompt Depth Anything preprocessing: the official ``load_image`` / ``load_depth`` (promptda/utils/io_wrapper.py @ |
| 9161547c) behind the transformers image-processor interface. That module also imports open3d / matplotlib for |
| visualisation, so its two loaders are re-typed here line by line (checked bit-exact against the originals). |
| |
| Image: RGB in [0, 1] (the model normalises it), longer side at most ``max_size`` (cv2 INTER_AREA, sides floored to |
| multiples of 14). A path is read with imageio, as the original. Prompt depth: meters, (h, w); a .png path is |
| millimeters (/ 1000), a .npz path holds "depth", as the original.""" |
|
|
| import cv2 |
| import imageio |
| import numpy as np |
| import torch |
| from PIL import Image |
| from transformers.feature_extraction_utils import BatchFeature |
| from transformers.image_processing_utils import BaseImageProcessor |
|
|
|
|
| def _to_tensor(arr): |
| if arr.ndim == 2: |
| arr = arr[:, :, np.newaxis] |
| return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) |
|
|
|
|
| def _ensure_multiple_of(x, multiple_of=14): |
| return int(x // multiple_of * multiple_of) |
|
|
|
|
| class PromptDepthAnythingImageProcessor(BaseImageProcessor): |
| model_input_names = ["pixel_values", "prompt_depth"] |
|
|
| def __init__(self, max_size: int = 1008, **kwargs): |
| super().__init__(**kwargs) |
| self.max_size = max_size |
|
|
| def load_image(self, image, max_size: int, multiple_of: int = 14): |
| |
| if isinstance(image, str): |
| image = np.asarray(imageio.imread(image)).astype(np.float32) |
| else: |
| image = np.asarray(image.convert("RGB") if isinstance(image, Image.Image) else image).astype(np.float32) |
| image = image / 255. |
|
|
| max_size = max_size // multiple_of * multiple_of |
| if max(image.shape) > max_size: |
| h, w = image.shape[:2] |
| scale = max_size / max(h, w) |
| tar_h = _ensure_multiple_of(h * scale) |
| tar_w = _ensure_multiple_of(w * scale) |
| image = cv2.resize(image, (tar_w, tar_h), interpolation=cv2.INTER_AREA) |
| return _to_tensor(image) |
|
|
| @staticmethod |
| def load_depth(depth): |
| |
| if isinstance(depth, str): |
| if depth.endswith('.png'): |
| depth = np.asarray(imageio.imread(depth)).astype(np.float32) |
| depth = depth / 1000. |
| elif depth.endswith('.npz'): |
| depth = np.load(depth)['depth'] |
| else: |
| raise ValueError(f"Unsupported depth format: {depth}") |
| elif isinstance(depth, torch.Tensor): |
| depth = depth.cpu().numpy() |
| return _to_tensor(np.asarray(depth)) |
|
|
| def preprocess(self, images, prompt_depth=None, max_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature: |
| """images: one image or a list (paths, PIL images or RGB uint8 arrays); prompt_depth: one depth map per image |
| (paths, or arrays in meters). pixel_values (B, 3, H, W) in [0, 1]; prompt_depth (B, 1, h, w).""" |
| if not isinstance(images, (list, tuple)): |
| images = [images] |
| tensors = [self.load_image(image, max_size or self.max_size)[0] for image in images] |
| if len({t.shape for t in tensors}) > 1: |
| raise ValueError("images resize to different shapes; process them one at a time") |
| data = {"pixel_values": torch.stack(tensors)} |
| if prompt_depth is not None: |
| if not isinstance(prompt_depth, (list, tuple)): |
| prompt_depth = [prompt_depth] |
| if len(prompt_depth) != len(images): |
| raise ValueError("give one prompt depth per image") |
| data["prompt_depth"] = torch.stack([self.load_depth(depth)[0] for depth in prompt_depth]) |
| return BatchFeature(data, tensor_type=return_tensors) |
|
|
|
|
| __all__ = ["PromptDepthAnythingImageProcessor"] |
|
|
