File size: 4,005 Bytes
d92a7a8 | 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 | """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): # io_wrapper.to_tensor_func
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): # io_wrapper.ensure_multiple_of
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):
# io_wrapper.load_image, also taking a PIL image / RGB uint8 array
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):
# io_wrapper.load_depth, also taking an array / tensor in meters
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"]
|