File size: 3,089 Bytes
653a5eb | 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 | """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: # infer.py
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
# body of infer.py _load_images over the given frames
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: # portrait -> landscape
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"]
|