File size: 3,220 Bytes
1ef9874 | 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 | """Image resize/pad and state-selection helpers used to build the pi0.5 training data.
The functions below are copied VERBATIM from the benchmark repo's
task1_isaacsim/baseline/preprocess.py (the rest of that file is the offline
dataset-conversion tool and is not needed at inference time). They are what
`reference/policy_adapter_sim.py` imports as `from preprocess import ...`.
Use this exact code for live images: bilinear + antialias resizing gives
slightly different pixels in OpenCV, PIL and torchvision, and the policy was
trained on the output of THIS implementation.
from preprocess import resize_with_pad
img224 = resize_with_pad(rgb_uint8_hwc) # (H, W, 3) uint8 RGB -> (224, 224, 3) uint8 RGB
Requires numpy and torch.
"""
from __future__ import annotations
import numpy as np
import torch
# Name-based state projections. "full" is a no-op (identity); "compact"
# selects each arm's measured joint positions + gripper joint value, in
# the same order as the franka_fr3_duo embodiment's ai_data_contract.
COMPACT_STATE_NAMES = (
[f"franka_robot_left_measured_joint_states_left_fr3v2_joint{i}" for i in range(1, 8)]
+ ["franka_robot_left_joint_states_robotiq_85_left_knuckle_joint"]
+ [f"franka_robot_right_measured_joint_states_right_fr3v2_joint{i}" for i in range(1, 8)]
+ ["franka_robot_right_joint_states_robotiq_85_left_knuckle_joint"]
)
def _even(n: int) -> int:
return n - (n % 2)
def fit_resolution(src: tuple[int, int], target: tuple[int, int]) -> tuple[int, int]:
"""Largest (w, h) that fits inside `target` while keeping `src`'s aspect ratio."""
sw, sh = src
tw, th = target
scale = min(tw / sw, th / sh)
return max(2, _even(round(sw * scale))), max(2, _even(round(sh * scale)))
def pad_to_size(img: np.ndarray, size: tuple[int, int]) -> np.ndarray:
"""Center `img` (H,W,3 uint8) on a black canvas of `size` = (W, H)."""
w, h = size
ih, iw = img.shape[:2]
if (iw, ih) == (w, h):
return img
canvas = np.zeros((h, w, 3), dtype=img.dtype)
y0, x0 = (h - ih) // 2, (w - iw) // 2
canvas[y0 : y0 + ih, x0 : x0 + iw] = img
return canvas
def _resize_to_uint8(img: torch.Tensor, width: int, height: int) -> np.ndarray:
"""(3,H,W) float32 in [0,1] - how LeRobotDataset decodes video - to (H,W,3) uint8."""
if (int(img.shape[2]), int(img.shape[1])) != (width, height):
img = torch.nn.functional.interpolate(
img.unsqueeze(0), size=(height, width), mode="bilinear", align_corners=False, antialias=True
)[0]
arr = (img.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8).permute(1, 2, 0).numpy()
return np.ascontiguousarray(arr)
# --- Added for the handoff (not in the original file): the same pipeline for live uint8 frames. ---
def resize_with_pad(img: np.ndarray, size: tuple[int, int] = (224, 224)) -> np.ndarray:
"""(H, W, 3) uint8 RGB -> (size[1], size[0], 3) uint8 RGB, exactly as the training images were made."""
h, w = img.shape[:2]
inner_w, inner_h = fit_resolution((w, h), size)
tensor = torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float() / 255.0
return pad_to_size(_resize_to_uint8(tensor, inner_w, inner_h), size)
|