mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-15 14:02:14 +00:00
fix(robocasa): render overlay text once
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user