code/scripts/render_oven_metric_frame.py

from collections import Counter
from pathlib import Path
import argparse
import json
import math
import os
import pickle
import sys
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd
from PIL import Image, ImageDraw, ImageFont


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from rr_label_study.oven_study import (
    BimanualTakeTrayOutOfOven,
    DEFAULT_PPRE_TAU,
    FULL_CAMERA_SET,
    MotionTemplates,
    ReplayCache,
    Shape,
    _apply_relative_pose,
    _camera_point_cloud,
    _camera_file,
    _extract_sequence_poses,
    _frame_metrics,
    _launch_replay_env,
    _load_demo,
    _load_mask,
    _pregrasp_corridor_rel_poses,
    _project_points,
    _sample_full_tray_points,
    _sample_grasp_points,
    _visibility_projection_details,
)


VISIBILITY_CAMERAS = ["front", "wrist_right", "wrist_left"]
PATH_CAMERAS = ["front", "overhead", "over_shoulder_left"]
ALL_METRICS_CAMERAS = ["front", "wrist_right", "wrist_left"]
RGB_SIZE = 128
PANEL_SIZE = 256
HEADER_HEIGHT = 72
FOOTER_HEIGHT = 210


def _font(size: int) -> ImageFont.FreeTypeFont:
    candidates = [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf",
    ]
    for path in candidates:
        if Path(path).exists():
            return ImageFont.truetype(path, size=size)
    return ImageFont.load_default()


FONT_SM = _font(14)
FONT_MD = _font(18)
FONT_LG = _font(24)


def _load_rgb(episode_dir: Path, camera_name: str, frame_index: int) -> Image.Image:
    path = _camera_file(episode_dir, camera_name, "rgb", frame_index)
    return Image.open(path).convert("RGB")


def _scale_image(image: Image.Image, size: int = PANEL_SIZE) -> Image.Image:
    return image.resize((size, size), resample=Image.Resampling.BILINEAR)


def _new_panel(background: Tuple[int, int, int] = (0, 0, 0)) -> Image.Image:
    return Image.new("RGB", (PANEL_SIZE, PANEL_SIZE), background)


def _draw_badge(draw: ImageDraw.ImageDraw, x: int, y: int, text: str, fill: Tuple[int, int, int]) -> None:
    left = x
    top = y
    width = 12 + int(len(text) * 8.5)
    height = 24
    draw.rounded_rectangle((left, top, left + width, top + height), radius=6, fill=fill)
    draw.text((left + 6, top + 4), text, fill=(255, 255, 255), font=FONT_SM)


def _draw_panel_title(image: Image.Image, title: str, subtitle: Optional[str] = None) -> None:
    draw = ImageDraw.Draw(image)
    draw.rectangle((0, 0, PANEL_SIZE, 30), fill=(0, 0, 0))
    draw.text((8, 6), title, fill=(255, 255, 255), font=FONT_MD)
    if subtitle:
        draw.rectangle((0, PANEL_SIZE - 26, PANEL_SIZE, PANEL_SIZE), fill=(0, 0, 0))
        draw.text((8, PANEL_SIZE - 22), subtitle, fill=(230, 230, 230), font=FONT_SM)


def _draw_point(draw: ImageDraw.ImageDraw, xy: Tuple[int, int], color: Tuple[int, int, int], radius: int = 3) -> None:
    x, y = xy
    draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=color)


def _draw_polyline(draw: ImageDraw.ImageDraw, coords: Sequence[Tuple[int, int]], color: Tuple[int, int, int], width: int = 2) -> None:
    if len(coords) >= 2:
        draw.line(coords, fill=color, width=width)


def _draw_segmented_polyline(
    draw: ImageDraw.ImageDraw,
    coords: Sequence[Tuple[int, int]],
    colors: Sequence[Tuple[int, int, int]],
    width: int = 3,
) -> None:
    if len(coords) < 2:
        return
    for idx in range(len(coords) - 1):
        color = colors[min(idx + 1, len(colors) - 1)]
        draw.line((coords[idx], coords[idx + 1]), fill=color, width=width)


def _bar(draw: ImageDraw.ImageDraw, x: int, y: int, w: int, h: int, value: float, label: str, color: Tuple[int, int, int], threshold: Optional[float] = None) -> None:
    draw.text((x, y - 18), f"{label}: {value:.3f}", fill=(255, 255, 255), font=FONT_SM)
    draw.rounded_rectangle((x, y, x + w, y + h), radius=5, outline=(120, 120, 120), fill=(20, 20, 20))
    filled_w = int(max(0.0, min(1.0, value)) * w)
    if filled_w > 0:
        if filled_w >= 12:
            draw.rounded_rectangle((x, y, x + filled_w, y + h), radius=5, fill=color)
        else:
            draw.rectangle((x, y, x + filled_w, y + h), fill=color)
    if threshold is not None:
        tx = x + int(max(0.0, min(1.0, threshold)) * w)
        draw.line((tx, y - 2, tx, y + h + 2), fill=(255, 255, 255), width=2)


def _project_coords(
    points_world: np.ndarray,
    point_cloud_world: np.ndarray,
    extrinsics: np.ndarray,
    intrinsics: np.ndarray,
) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
    return _visibility_projection_details(
        points_world, point_cloud_world, extrinsics, intrinsics
    )


def _infer_mask_handle(mask: np.ndarray, coords: Sequence[Tuple[int, int]]) -> Optional[int]:
    values = [int(mask[y, x]) for x, y in coords if int(mask[y, x]) != 0]
    if not values:
        return None
    return Counter(values).most_common(1)[0][0]


def _masked_rgb(rgb: Image.Image, binary_mask: np.ndarray) -> Image.Image:
    rgb_np = np.asarray(rgb).copy()
    out = np.zeros_like(rgb_np)
    out[binary_mask] = rgb_np[binary_mask]
    return Image.fromarray(out, mode="RGB")


def _scaled_coords(coords: Iterable[Tuple[int, int]], scale: float) -> List[Tuple[int, int]]:
    return [(int(round(x * scale)), int(round(y * scale))) for x, y in coords]


def _demo_gripper_pose(obs, arm_name: str) -> np.ndarray:
    side = getattr(obs, arm_name)
    return np.asarray(side.gripper_pose, dtype=np.float64)


def _world_pose_to_coords(
    pose_xyz: Sequence[float],
    extrinsics: np.ndarray,
    intrinsics: np.ndarray,
) -> Optional[Tuple[int, int]]:
    points = np.asarray([pose_xyz], dtype=np.float64)
    uv, camera_xyz = _project_points(points, extrinsics, intrinsics)
    u, v = uv[0]
    depth = camera_xyz[0, 2]
    if depth <= 0 or not (0 <= u < RGB_SIZE and 0 <= v < RGB_SIZE):
        return None
    return int(round(float(u))), int(round(float(v)))


def _compose_visibility_panel(
    episode_dir: Path,
    demo,
    frame_index: int,
    camera_name: str,
    tray_pose: np.ndarray,
    left_pose: np.ndarray,
    right_pose: np.ndarray,
    templates: MotionTemplates,
    per_camera_vis: float,
) -> Image.Image:
    rgb = _load_rgb(episode_dir, camera_name, frame_index)
    point_cloud, extrinsics, intrinsics = _camera_point_cloud(
        episode_dir, demo, frame_index, camera_name
    )

    tray_points = _sample_full_tray_points(tray_pose)
    grasp_points = _sample_grasp_points(templates, tray_pose)
    tray_proj, tray_visible = _project_coords(
        tray_points, point_cloud, extrinsics, intrinsics
    )
    grasp_proj, grasp_visible = _project_coords(
        grasp_points, point_cloud, extrinsics, intrinsics
    )

    scene_panel = _scale_image(rgb)
    _draw_panel_title(
        scene_panel,
        f"{camera_name} scene",
        f"metric vis={per_camera_vis:.3f} | visible={len(grasp_visible)}/{len(grasp_proj) or 1}",
    )
    draw = ImageDraw.Draw(scene_panel)
    for xy in _scaled_coords(tray_proj, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (70, 170, 255), radius=2)
    for xy in _scaled_coords(grasp_proj, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (255, 80, 210), radius=3)
    for xy in _scaled_coords(grasp_visible, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (70, 255, 140), radius=3)
    _draw_badge(draw, 8, 34, "blue=tray", (40, 105, 180))
    _draw_badge(draw, 112, 34, "magenta=grasp", (145, 40, 140))
    _draw_badge(draw, 8, 64, "green=visible", (32, 130, 64))

    xray_panel = _new_panel()
    _draw_panel_title(xray_panel, f"{camera_name} x-ray", "blue=tray, magenta=grasp region")
    draw = ImageDraw.Draw(xray_panel)
    for xy in _scaled_coords(tray_proj, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (70, 170, 255), radius=2)
    for xy in _scaled_coords(grasp_proj, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (255, 80, 210), radius=3)
    left_xy = _world_pose_to_coords(left_pose[:3], extrinsics, intrinsics)
    right_xy = _world_pose_to_coords(right_pose[:3], extrinsics, intrinsics)
    if left_xy is not None:
        _draw_point(draw, _scaled_coords([left_xy], PANEL_SIZE / RGB_SIZE)[0], (255, 230, 70), radius=4)
    if right_xy is not None:
        _draw_point(draw, _scaled_coords([right_xy], PANEL_SIZE / RGB_SIZE)[0], (140, 110, 255), radius=4)

    visible_pixels = np.zeros((RGB_SIZE, RGB_SIZE), dtype=bool)
    for x, y in tray_visible + grasp_visible:
        visible_pixels[y, x] = True
    masked_panel = _masked_rgb(rgb, visible_pixels)
    masked_panel = _scale_image(masked_panel)
    _draw_panel_title(
        masked_panel,
        f"{camera_name} depth-visible",
        f"visible grasp pts={len(grasp_visible)}/{len(grasp_proj) or 1}",
    )
    draw = ImageDraw.Draw(masked_panel)
    for xy in _scaled_coords(grasp_proj, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (180, 80, 160), radius=2)
    for xy in _scaled_coords(grasp_visible, PANEL_SIZE / RGB_SIZE):
        _draw_point(draw, xy, (70, 255, 140), radius=3)

    row = Image.new("RGB", (PANEL_SIZE * 3, PANEL_SIZE), (15, 15, 15))
    row.paste(scene_panel, (0, 0))
    row.paste(xray_panel, (PANEL_SIZE, 0))
    row.paste(masked_panel, (PANEL_SIZE * 2, 0))
    return row


def _compose_visibility_focus(
    episode_dir: Path,
    demo,
    frame_index: int,
    tray_pose: np.ndarray,
    left_pose: np.ndarray,
    right_pose: np.ndarray,
    templates: MotionTemplates,
    frame_row: pd.Series,
) -> Image.Image:
    canvas = Image.new("RGB", (PANEL_SIZE * 3, HEADER_HEIGHT + PANEL_SIZE * len(VISIBILITY_CAMERAS)), (12, 12, 12))
    draw = ImageDraw.Draw(canvas)
    draw.rectangle((0, 0, canvas.size[0], HEADER_HEIGHT), fill=(0, 0, 0))
    draw.text((12, 10), f"Visibility Metric Debug | frame {frame_index}", fill=(255, 255, 255), font=FONT_LG)
    draw.text(
        (12, 40),
        f"three_view_visibility={frame_row['three_view_visibility']:.3f} | full_view_visibility={frame_row['full_view_visibility']:.3f} | y_ready={int(frame_row['y_ready'])} | phase_score={float(frame_row.get('phase_score', frame_row['phase_switch'])):.3f}",
        fill=(235, 235, 235),
        font=FONT_SM,
    )
    for row_index, camera_name in enumerate(VISIBILITY_CAMERAS):
        per_camera_vis = float(_camera_grasp_visibility(episode_dir, demo, frame_index, camera_name, tray_pose, templates))
        row = _compose_visibility_panel(
            episode_dir=episode_dir,
            demo=demo,
            frame_index=frame_index,
            camera_name=camera_name,
            tray_pose=tray_pose,
            left_pose=left_pose,
            right_pose=right_pose,
            templates=templates,
            per_camera_vis=per_camera_vis,
        )
        canvas.paste(row, (0, HEADER_HEIGHT + row_index * PANEL_SIZE))
    return canvas


def _camera_grasp_visibility(
    episode_dir: Path,
    demo,
    frame_index: int,
    camera_name: str,
    tray_pose: np.ndarray,
    templates: MotionTemplates,
) -> float:
    point_cloud, extrinsics, intrinsics = _camera_point_cloud(
        episode_dir, demo, frame_index, camera_name
    )
    grasp_points = _sample_grasp_points(templates, tray_pose)
    proj, vis = _project_coords(grasp_points, point_cloud, extrinsics, intrinsics)
    return float(len(vis) / len(proj)) if proj else 0.0


def _compose_path_quality(
    episode_dir: Path,
    demo,
    frame_index: int,
    tray_pose: np.ndarray,
    task_base_pose: np.ndarray,
    left_pose: np.ndarray,
    right_pose: np.ndarray,
    templates: MotionTemplates,
    frame_row: pd.Series,
    left_grasped: bool,
) -> Image.Image:
    canvas = Image.new("RGB", (PANEL_SIZE * 2, HEADER_HEIGHT + PANEL_SIZE * 2), (12, 12, 12))
    draw = ImageDraw.Draw(canvas)
    draw.rectangle((0, 0, canvas.size[0], HEADER_HEIGHT), fill=(0, 0, 0))
    draw.text((12, 10), f"Path Quality Debug | frame {frame_index}", fill=(255, 255, 255), font=FONT_LG)
    draw.text(
        (12, 40),
        f"p_pre={frame_row['p_pre']:.3f} y_pre={int(frame_row['y_pre'])} | p_ext={frame_row['p_ext']:.3f} y_ext={int(frame_row['y_ext'])} | grasped={int(left_grasped)}",
        fill=(235, 235, 235),
        font=FONT_SM,
    )

    pregrasp_corridor = [
        _apply_relative_pose(tray_pose, rel_pose)
        for rel_pose in _pregrasp_corridor_rel_poses(templates)
    ]
    extract_poses = _extract_sequence_poses(tray_pose, task_base_pose, templates)
    plan_poses = [*pregrasp_corridor, extract_poses[1], *extract_poses[2:5]]
    colors = ([(255, 220, 0)] * len(pregrasp_corridor)) + [
        (255, 140, 0),
        (80, 255, 120),
        (80, 255, 120),
        (80, 255, 120),
    ]

    demo_left_trail = np.asarray(
        [_demo_gripper_pose(demo[i], "left")[:3] for i in range(max(0, frame_index - 12), frame_index + 1)],
        dtype=np.float64,
    )
    demo_right_trail = np.asarray(
        [_demo_gripper_pose(demo[i], "right")[:3] for i in range(max(0, frame_index - 12), frame_index + 1)],
        dtype=np.float64,
    )

    for camera_slot, camera_name in enumerate(PATH_CAMERAS):
        panel = _scale_image(_load_rgb(episode_dir, camera_name, frame_index))
        _draw_panel_title(panel, camera_name, "demo trails + planned extract path")
        draw = ImageDraw.Draw(panel)
        extrinsics = demo[frame_index].misc[f"{camera_name}_camera_extrinsics"]
        intrinsics = demo[frame_index].misc[f"{camera_name}_camera_intrinsics"]
        left_trail_xy = _project_world_polyline(demo_left_trail, extrinsics, intrinsics)
        right_trail_xy = _project_world_polyline(demo_right_trail, extrinsics, intrinsics)
        plan_xy = [
            xy
            for pose in plan_poses
            for xy in [_world_pose_to_coords(pose[:3], extrinsics, intrinsics)]
            if xy is not None
        ]
        _draw_polyline(draw, _scaled_coords(left_trail_xy, PANEL_SIZE / RGB_SIZE), (80, 220, 255), width=3)
        _draw_polyline(draw, _scaled_coords(right_trail_xy, PANEL_SIZE / RGB_SIZE), (150, 110, 255), width=3)
        _draw_segmented_polyline(draw, _scaled_coords(plan_xy, PANEL_SIZE / RGB_SIZE), colors, width=4)
        current_left_xy = _world_pose_to_coords(left_pose[:3], extrinsics, intrinsics)
        current_right_xy = _world_pose_to_coords(right_pose[:3], extrinsics, intrinsics)
        if current_left_xy is not None:
            _draw_point(draw, _scaled_coords([current_left_xy], PANEL_SIZE / RGB_SIZE)[0], (255, 255, 255), radius=5)
        if current_right_xy is not None:
            _draw_point(draw, _scaled_coords([current_right_xy], PANEL_SIZE / RGB_SIZE)[0], (255, 255, 255), radius=5)
        for pose, color in zip(plan_poses, colors):
            xy = _world_pose_to_coords(pose[:3], extrinsics, intrinsics)
            if xy is not None:
                _draw_point(draw, _scaled_coords([xy], PANEL_SIZE / RGB_SIZE)[0], color, radius=4)
        x = (camera_slot % 2) * PANEL_SIZE
        y = HEADER_HEIGHT + (camera_slot // 2) * PANEL_SIZE
        canvas.paste(panel, (x, y))

    text_panel = Image.new("RGB", (PANEL_SIZE, PANEL_SIZE), (18, 18, 18))
    draw = ImageDraw.Draw(text_panel)
    draw.text((12, 12), "Metric State", fill=(255, 255, 255), font=FONT_LG)
    _bar(draw, 12, 56, PANEL_SIZE - 24, 18, float(frame_row["p_pre"]), "p_pre", (255, 220, 0), threshold=DEFAULT_PPRE_TAU)
    _bar(draw, 12, 108, PANEL_SIZE - 24, 18, float(frame_row["p_ext"]), "p_ext", (80, 255, 120), threshold=0.45)
    draw.text((12, 146), f"y_pre={int(frame_row['y_pre'])} y_ext={int(frame_row['y_ext'])}", fill=(255, 255, 255), font=FONT_MD)
    draw.text((12, 172), f"pre_prog={float(frame_row.get('pregrasp_progress', 0.0)):.3f}", fill=(255, 255, 255), font=FONT_MD)
    draw.text((12, 198), f"left_grasped={int(left_grasped)} door={float(frame_row['door_angle']):.3f}", fill=(255, 255, 255), font=FONT_MD)
    draw.text((12, 224), f"y_retrieve={int(frame_row.get('y_retrieve', 0))} phase={int(frame_row['phase_switch'])} score={float(frame_row.get('phase_score', frame_row['phase_switch'])):.3f}", fill=(255, 255, 255), font=FONT_MD)
    canvas.paste(text_panel, (PANEL_SIZE, HEADER_HEIGHT + PANEL_SIZE))
    return canvas


def _project_world_polyline(points_world: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray) -> List[Tuple[int, int]]:
    if len(points_world) == 0:
        return []
    uv, camera_xyz = _project_points(points_world, extrinsics, intrinsics)
    coords = []
    for (u, v), (_, _, depth) in zip(uv, camera_xyz):
        if depth <= 0 or not (0 <= u < RGB_SIZE and 0 <= v < RGB_SIZE):
            continue
        coords.append((int(round(float(u))), int(round(float(v)))))
    return coords


def _compose_all_metrics(
    episode_name: str,
    episode_dir: Path,
    demo,
    frame_index: int,
    frame_row: pd.Series,
) -> Image.Image:
    banner_color = (145, 30, 30) if int(frame_row["phase_switch"]) == 0 else (25, 115, 40)
    canvas = Image.new("RGB", (PANEL_SIZE * 3, PANEL_SIZE + FOOTER_HEIGHT), (12, 12, 12))
    draw = ImageDraw.Draw(canvas)
    draw.rectangle((0, 0, canvas.size[0], 36), fill=banner_color)
    draw.text(
        (12, 7),
        f"{episode_name} | frame {frame_index:03d} | {'REVEAL' if int(frame_row['phase_switch']) == 0 else 'RETRIEVE'}",
        fill=(255, 255, 255),
        font=FONT_LG,
    )
    for idx, camera_name in enumerate(ALL_METRICS_CAMERAS):
        panel = _scale_image(_load_rgb(episode_dir, camera_name, frame_index))
        _draw_panel_title(panel, camera_name, None)
        canvas.paste(panel, (idx * PANEL_SIZE, 36))

    footer_y = 36 + PANEL_SIZE
    draw.rectangle((0, footer_y, canvas.size[0], canvas.size[1]), fill=(18, 18, 18))
    _bar(draw, 16, footer_y + 34, 220, 18, float(frame_row["three_view_visibility"]), "three_view_visibility", (60, 160, 255), threshold=0.35)
    _bar(draw, 272, footer_y + 34, 220, 18, float(frame_row["full_view_visibility"]), "full_view_visibility", (80, 190, 255), threshold=0.35)
    _bar(draw, 528, footer_y + 34, 220, 18, float(frame_row["p_pre"]), "p_pre", (255, 220, 0), threshold=DEFAULT_PPRE_TAU)
    _bar(draw, 16, footer_y + 96, 220, 18, float(frame_row["p_ext"]), "p_ext", (80, 255, 120), threshold=0.45)
    draw.text((272, footer_y + 84), f"y_pre={int(frame_row['y_pre'])} y_ext={int(frame_row['y_ext'])} pre_prog={float(frame_row.get('pregrasp_progress', 0.0)):.3f}", fill=(255, 255, 255), font=FONT_MD)
    draw.text((272, footer_y + 108), f"y_ready={int(frame_row['y_ready'])} phase_switch={int(frame_row['phase_switch'])} phase_score={float(frame_row.get('phase_score', frame_row['phase_switch'])):.3f}", fill=(255, 255, 255), font=FONT_MD)
    draw.text((528, footer_y + 84), f"door_angle={float(frame_row['door_angle']):.3f}", fill=(255, 255, 255), font=FONT_MD)
    draw.text((528, footer_y + 108), f"left_open={float(frame_row['left_gripper_open']):.1f} right_open={float(frame_row['right_gripper_open']):.1f}", fill=(255, 255, 255), font=FONT_MD)
    return canvas


def _save_if_requested(image: Image.Image, path: Optional[Path]) -> None:
    if path is None:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    image.save(path)


def _load_templates(path: Path) -> MotionTemplates:
    if path.suffix.lower() == ".json":
        with path.open("r", encoding="utf-8") as handle:
            payload = json.load(handle)
        return MotionTemplates.from_json(payload["templates"])
    with path.open("rb") as handle:
        return pickle.load(handle)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--episode-dir", required=True)
    parser.add_argument("--templates-pkl", required=True)
    parser.add_argument("--dense-csv", required=True)
    parser.add_argument("--frame-index", type=int, required=True)
    parser.add_argument("--checkpoint-stride", type=int, default=16)
    parser.add_argument("--visibility-out")
    parser.add_argument("--path-out")
    parser.add_argument("--all-out")
    args = parser.parse_args()

    episode_dir = Path(args.episode_dir)
    templates = _load_templates(Path(args.templates_pkl))
    frame_df = pd.read_csv(args.dense_csv)
    if "frame_index" in frame_df.columns:
        matches = frame_df.loc[frame_df["frame_index"] == int(args.frame_index)]
        if len(matches) == 0:
            raise ValueError(f"frame_index {args.frame_index} not found in {args.dense_csv}")
        frame_row = matches.iloc[0]
    else:
        frame_row = frame_df.iloc[int(args.frame_index)]
    demo = _load_demo(episode_dir)

    env = _launch_replay_env()
    try:
        task = env.get_task(BimanualTakeTrayOutOfOven)
        cache = ReplayCache(task, demo, checkpoint_stride=args.checkpoint_stride)
        cache.reset()
        cache.step_to(int(args.frame_index))
        state = cache.current_state()
        left_grasped = any(
            obj.get_name() == "tray"
            for obj in task._scene.robot.left_gripper.get_grasped_objects()
        )

        visibility_out = Path(args.visibility_out) if args.visibility_out else None
        path_out = Path(args.path_out) if args.path_out else None
        all_out = Path(args.all_out) if args.all_out else None

        if visibility_out is not None:
            visibility_image = _compose_visibility_focus(
                episode_dir=episode_dir,
                demo=demo,
                frame_index=int(args.frame_index),
                tray_pose=state.tray_pose,
                left_pose=state.left_gripper_pose,
                right_pose=state.right_gripper_pose,
                templates=templates,
                frame_row=frame_row,
            )
            _save_if_requested(visibility_image, visibility_out)

        if path_out is not None:
            path_image = _compose_path_quality(
                episode_dir=episode_dir,
                demo=demo,
                frame_index=int(args.frame_index),
                tray_pose=state.tray_pose,
                task_base_pose=task._task.get_base().get_pose(),
                left_pose=state.left_gripper_pose,
                right_pose=state.right_gripper_pose,
                templates=templates,
                frame_row=frame_row,
                left_grasped=left_grasped,
            )
            _save_if_requested(path_image, path_out)

        if all_out is not None:
            all_image = _compose_all_metrics(
                episode_name=episode_dir.name,
                episode_dir=episode_dir,
                demo=demo,
                frame_index=int(args.frame_index),
                frame_row=frame_row,
            )
            _save_if_requested(all_image, all_out)
    finally:
        env.shutdown()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())