| """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 |
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| 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) |
|
|
