From 519ec054373992e507bec8a574f66cc828dc7cc5 Mon Sep 17 00:00:00 2001 From: wayrise Date: Wed, 9 Sep 2026 02:24:29 +0800 Subject: [PATCH] fix(robodojo): decode JPEG frames with the recorder's channel order RoboDojo writes its camera buffers by handing an RGB array straight to cv2.imencode, which expects BGR, so R and B are swapped inside the file. The reader decoded them with PIL, which returns that stored order, so training saw channel-swapped frames. Deploy is unaffected: obs_preprocess decodes whatever the eval client sends, and every in-repo client encodes true RGB, so the model trained on swapped frames and evaluated on correct ones. robotwin.py documents and handles exactly this for the RoboTwin corpus. Confirmed on real data rather than by reading: each corpus ships preview_video/*.mp4 produced by the recorder itself, and over saturated pixels the reader now matches those instead of inverting them. sim corpus, production config |preview - reader| 40..57 -> 3..16 real corpus, no colour jitter |preview - reader| 31.74 -> 4.92 The residual is the JPEG vs H.264 difference; the sim spread is the configured colour jitter. Per-channel means went from transposed to aligned. The test fixture encoded with PIL, i.e. the opposite convention to the real recorder, which is why a wrong decoder passed: it now encodes the way RoboDojo does, so the existing three-camera channel assertions pin the real contract, and a new test states the property directly. Reverting the decoder fails 4 tests. Note for anyone with RoboDojo checkpoints: this changes the training input distribution, so weights trained before it were fitted on swapped frames. Closes #13. Co-Authored-By: Claude Opus 5 --- openwam/dataloader/robodojo.py | 26 ++++++++++++++++----- tests/test_robodojo_dataloader.py | 39 +++++++++++++++++++++++++++---- 2 files changed, 55 insertions(+), 10 deletions(-) diff --git a/openwam/dataloader/robodojo.py b/openwam/dataloader/robodojo.py index 5ee79443..78e160db 100644 --- a/openwam/dataloader/robodojo.py +++ b/openwam/dataloader/robodojo.py @@ -22,7 +22,6 @@ import bisect import hashlib -import io import json import os import time @@ -30,6 +29,7 @@ from pathlib import Path from typing import Any +import cv2 import h5py import numpy as np import torch @@ -357,14 +357,28 @@ def _jpeg_bytes(value: Any) -> bytes: def _decode_jpeg(value: Any, *, source: str) -> Image.Image: + """Decode JPEG bytes from HDF5 to a PIL RGB image. + + RoboDojo encodes frames by passing RGB arrays directly to ``cv2.imencode`` + (which expects BGR), so R and B channels are swapped inside the JPEG. Using + ``cv2.imdecode`` reverses this swap, giving back the original RGB order — no + further conversion needed. This is the same convention ``robotwin.py`` + documents for the RoboTwin corpus. + + Decoding with PIL instead returns the stored order, i.e. R and B swapped, + which is what this reader did until this was caught: the model trained on + channel-swapped frames while the deploy path (which decodes whatever the + eval client encodes, and every in-repo client encodes true RGB) served + correct ones. Confirmed against the corpora's own ``preview_video/*.mp4``, + which match this decoder and not PIL's, on both the sim and real corpora. + """ encoded = _jpeg_bytes(value) if not encoded: raise ValueError(f"{source}: JPEG entry is empty") - try: - with Image.open(io.BytesIO(encoded)) as image: - return image.convert("RGB").copy() - except Exception as error: - raise ValueError(f"{source}: could not decode JPEG") from error + array = cv2.imdecode(np.frombuffer(encoded, np.uint8), cv2.IMREAD_COLOR) + if array is None: + raise ValueError(f"{source}: could not decode JPEG") + return Image.fromarray(array) def discover_robodojo_tasks( diff --git a/tests/test_robodojo_dataloader.py b/tests/test_robodojo_dataloader.py index 583bb5e6..773142d0 100644 --- a/tests/test_robodojo_dataloader.py +++ b/tests/test_robodojo_dataloader.py @@ -5,6 +5,7 @@ import io from pathlib import Path +import cv2 import h5py import numpy as np import pytest @@ -59,10 +60,18 @@ def write_calibration(root: Path) -> Path: def encode_jpeg(color: tuple[int, int, int], height: int = 18, width: int = 20) -> bytes: - image = Image.new("RGB", (width, height), color) - buffer = io.BytesIO() - image.save(buffer, format="JPEG", quality=100, subsampling=0) - return buffer.getvalue() + """Encode ``color`` (RGB) the way the RoboDojo recorder does. + + It hands an RGB array straight to ``cv2.imencode``, which expects BGR, so R + and B end up swapped inside the file. Fixtures have to reproduce that or the + reader's ``cv2.imdecode`` — which exists to undo it — looks wrong under test + while being right on real data. + """ + rgb = np.zeros((height, width, 3), dtype=np.uint8) + rgb[:, :] = color + ok, buffer = cv2.imencode(".jpg", rgb, [int(cv2.IMWRITE_JPEG_QUALITY), 100]) + assert ok, "cv2.imencode failed" + return buffer.tobytes() def source_arrays(T: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: @@ -856,3 +865,25 @@ def test_registry_constructs_robodojo_and_exports_classes(tmp_path: Path): assert isinstance(dataset, MultiTaskRoboDojoDataset) assert isinstance(dataset._sub_datasets[0], RoboDojoDataset) assert dataset[0]["action"].shape == (2, 80) + + +def test_decode_jpeg_undoes_the_recorder_channel_swap(): + """The reader must return true RGB, not the stored order. + + RoboDojo writes its JPEGs with ``cv2.imencode`` fed an RGB array, so the file + holds R and B swapped. Decoding with PIL returns that stored order and feeds + the model channel-swapped frames, which is what this reader did until it was + caught against the corpora's own preview videos. Pinned here because CI never + sees a real episode. + """ + from openwam.dataloader.robodojo import _decode_jpeg + + red = (230, 20, 15) + frame = np.asarray(_decode_jpeg(encode_jpeg(red), source="test")) + + assert frame[0, 0, 0] > 200, "red must come back in channel 0" + assert frame[0, 0, 2] < 60, "blue must not have been swapped into channel 0" + + # And the PIL reading of the very same bytes is the swap this guards against. + stored = np.asarray(Image.open(io.BytesIO(encode_jpeg(red))).convert("RGB")) + assert stored[0, 0, 2] > 200 and stored[0, 0, 0] < 60