fix(robocasa): render overlay text once

This commit is contained in:
Pepijn
2026-07-15 12:07:23 +02:00
parent dbb7f5b769
commit 1f00078cc7
2 changed files with 69 additions and 15 deletions
+25 -15
View File
@@ -47,9 +47,7 @@ def _label_panel(img: np.ndarray, label: str) -> np.ndarray:
return img
def _overlay_text(
frame: np.ndarray, task: str | None, subtask: str | None, memory: str | None
) -> np.ndarray:
def _overlay_text(frame: np.ndarray, task: str | None, subtask: str | None, memory: str | None) -> np.ndarray:
"""Draw task / subtask / memory lines onto an (H, W, 3) uint8 frame.
Best-effort: returns the frame unchanged if OpenCV is unavailable.
@@ -59,31 +57,41 @@ def _overlay_text(
except ImportError:
return frame
lines = [f"{label}: {val}" for label, val in
(("Task", task), ("Subtask", subtask), ("Memory", memory)) if val]
lines = [
f"{label}: {val}" for label, val in (("Task", task), ("Subtask", subtask), ("Memory", memory)) if val
]
if not lines:
return frame
img = np.ascontiguousarray(frame)
img = np.ascontiguousarray(frame).copy()
font, scale, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 6
max_width = img.shape[1] - 2 * margin
y = 18
wrapped_lines: list[str] = []
for text in lines:
# naive width-based wrap so long memory strings stay on-frame
words, cur = text.split(), ""
wrapped: list[str] = []
for w in words:
cand = f"{cur} {w}".strip()
if cv2.getTextSize(cand, font, scale, 1)[0][0] > max_width and cur:
wrapped.append(cur)
wrapped_lines.append(cur)
cur = w
else:
cur = cand
wrapped.append(cur)
for line in wrapped:
cv2.putText(img, line, (margin, y), font, scale, (0, 0, 0), 3, cv2.LINE_AA)
cv2.putText(img, line, (margin, y), font, scale, (255, 255, 255), 1, cv2.LINE_AA)
y += 20
wrapped_lines.append(cur)
# Use a dark translucent header for contrast, then draw each label once.
# The previous black-outline + white-foreground technique rendered every
# glyph twice and looked like offset duplicate text in the MJPEG stream.
line_height = 20
header_height = min(img.shape[0], len(wrapped_lines) * line_height + 6)
backdrop = img.copy()
cv2.rectangle(backdrop, (0, 0), (img.shape[1], header_height), (0, 0, 0), -1)
cv2.addWeighted(backdrop, 0.55, img, 0.45, 0, dst=img)
y = 18
for line in wrapped_lines:
cv2.putText(img, line, (margin, y), font, scale, (255, 255, 255), 1, cv2.LINE_AA)
y += line_height
return img
@@ -196,7 +204,9 @@ def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) ->
pass
try:
server = ThreadingHTTPServer(("0.0.0.0", port), _Handler)
# Bind all interfaces intentionally so the viewer remains reachable
# through the documented SSH port-forwarding workflow.
server = ThreadingHTTPServer(("0.0.0.0", port), _Handler) # nosec B104
except OSError as exc:
logger.warning("[sim] could not start live stream on port %d: %s", port, exc)
print(f"[runtime] WARNING: live stream port {port} unavailable ({exc})", flush=True)
+44
View File
@@ -0,0 +1,44 @@
import sys
from types import SimpleNamespace
import numpy as np
from lerobot.runtime.sim_robocasa import _overlay_text
def test_overlay_draws_each_label_once(monkeypatch):
put_text_calls = []
rectangle_calls = []
def put_text(image, text, origin, font, scale, color, thickness, line_type):
put_text_calls.append((text, color, thickness))
return image
def rectangle(image, start, end, color, thickness):
rectangle_calls.append((start, end, color, thickness))
return image
def add_weighted(src1, alpha, src2, beta, gamma, *, dst):
dst[:] = src1 * alpha + src2 * beta + gamma
return dst
fake_cv2 = SimpleNamespace(
FONT_HERSHEY_SIMPLEX=0,
LINE_AA=16,
getTextSize=lambda text, font, scale, thickness: ((len(text) * 7, 10), 0),
putText=put_text,
rectangle=rectangle,
addWeighted=add_weighted,
)
monkeypatch.setitem(sys.modules, "cv2", fake_cv2)
frame = np.full((120, 480, 3), 200, dtype=np.uint8)
annotated = _overlay_text(frame, "close the fridge", "reach for the handle", None)
assert [call[0] for call in put_text_calls] == [
"Task: close the fridge",
"Subtask: reach for the handle",
]
assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
assert len(rectangle_calls) == 1
assert not np.shares_memory(annotated, frame)