| """Trace Anything preprocessing: scripts/infer.py ``_load_images`` (ByteDance-Seed/TraceAnything @ 54677b5e) behind the |
| transformers image-processor interface: portrait frames rotated to landscape, longer side resized to 512 (PIL bilinear), |
| every frame cropped to the first frame's size floored to multiples of 16, normalised to [-1, 1]; frame i of N gets the |
| time step i / (N - 1). infer.py reads files with cv2 (BGR -> RGB); a path is read the same way here, a PIL image / RGB |
| uint8 array is used as RGB. (infer.py also keeps every (N // 39)-th frame of scenes with more than 40 frames; that |
| subsampling is left to the caller.)""" |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| import torchvision.transforms as tvf |
| from PIL import Image |
| from transformers.feature_extraction_utils import BatchFeature |
| from transformers.image_processing_utils import BaseImageProcessor |
|
|
|
|
| def _resize_long_side(pil: Image.Image, long: int = 512) -> Image.Image: |
| w, h = pil.size |
| if w >= h: |
| return pil.resize((long, int(h * long / w)), Image.BILINEAR) |
| else: |
| return pil.resize((int(w * long / h), long), Image.BILINEAR) |
|
|
|
|
| class TraceAnythingImageProcessor(BaseImageProcessor): |
| model_input_names = ["pixel_values", "time_steps"] |
|
|
| def __init__(self, image_size: int = 512, **kwargs): |
| super().__init__(**kwargs) |
| self.image_size = image_size |
|
|
| def preprocess(self, images, image_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature: |
| """images: the frames of one video, in order (paths, PIL images or RGB uint8 arrays). pixel_values |
| (1, N, 3, H, W) in [-1, 1]; time_steps (1, N) float64 (i / (N - 1), 0 for one frame).""" |
| long = image_size or self.image_size |
| |
| tfm = tvf.Compose([tvf.ToTensor(), tvf.Normalize((0.5,)*3, (0.5,)*3)]) |
| frames, times, target = [], [], None |
| for i, image in enumerate(images): |
| if isinstance(image, str): |
| arr = cv2.imread(image) |
| if arr is None: |
| raise FileNotFoundError(f"Failed to read {image}") |
| pil = Image.fromarray(cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)) |
| else: |
| pil = image.convert("RGB") if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image)) |
|
|
| W0, H0 = pil.size |
| if H0 > W0: |
| pil = pil.transpose(Image.Transpose.ROTATE_90) |
|
|
| pil = _resize_long_side(pil, long) |
| if target is None: |
| H, W = pil.size[1], pil.size[0] |
| target = (H - H % 16, W - W % 16) |
| Ht, Wt = target |
| pil = pil.crop((0, 0, Wt, Ht)) |
|
|
| frames.append(tfm(pil)) |
| times.append(i / (len(images) - 1) if len(images) > 1 else 0.0) |
| data = {"pixel_values": torch.stack(frames)[None], "time_steps": torch.tensor([times], dtype=torch.float64)} |
| return BatchFeature(data, tensor_type=return_tensors) |
|
|
|
|
| __all__ = ["TraceAnythingImageProcessor"] |
|
|
