质检

"""Check (recording): the share of frames where a camera had a gap.

In:  the canonical recording (the frames of every camera and which of them are valid).
Out: one check result, "pass" or "warn", with the worst share and where the first gap is.
Why: a camera that drops frames (a USB hiccup on a head-mounted rig, a stalled robot camera) breaks
     video-based training; the check tells you before the recording reaches a dataset version.
"""

from dataclasses import dataclass

from embodiflow import Check, CheckWriter, Context, Recording, algorithm, param


@dataclass
class Params:
    max_ratio: float = param(0.01, minimum=0.0, maximum=1.0, title="Largest accepted gap share")


@algorithm.check(
    name="frame-gaps",
    title="Frame gaps",
    checks=[Check("frame_gap_ratio", "Share of frames missing from a camera")],
    params=Params,
)
def run(recording: Recording, out: CheckWriter, ctx: Context[Params]) -> None:
    worst, worst_key, first_gap = 0.0, "", None
    for camera in recording.cameras:
        valid = camera.valid()
        ratio = float(1.0 - valid.mean()) if len(valid) else 0.0
        if ratio > worst:
            worst, worst_key = ratio, camera.key
            first_gap = int((~valid).argmax())
    status = "pass" if worst <= ctx.params.max_ratio else "warn"
    at = {"startS": first_gap / recording.fps, "endS": (first_gap + 1) / recording.fps, "camera": worst_key} if first_gap is not None else None
    out.result("frame_gap_ratio", status, value=worst, threshold=ctx.params.max_ratio, message=f"{worst:.1%} of {worst_key or 'camera'} frames are gaps", at=at)

"""Check (raw): the MCAP file has the topics a standardizer needs, at a sane rate.

In:  the uploaded files (raw), before any standardizing: an MCAP file recorded from ROS 2.
Out: two check results: all required topics are there, and each publishes fast enough.
Why: a missing or slow camera topic is cheapest to find right after the upload, not after a long workflow.
     A missing topic fails the check but not the step; an upload without an MCAP is "not applicable".
"""

from dataclasses import dataclass

from embodiflow import Check, CheckWriter, Context, RawFiles, algorithm, param


@dataclass
class Params:
    required: list[str] = param(["/camera/image"], widget="tags", min_items=1, title="Required topics")
    min_rate_hz: float = param(5.0, minimum=0.0, title="Lowest accepted rate (Hz)")


@algorithm.check(
    name="raw-topics",
    title="Required MCAP topics",
    reads="raw",
    checks=[Check("topics_present", "Required topics are present"), Check("topic_rate", "Required topics publish fast enough")],
    params=Params,
)
def run(raw: RawFiles, out: CheckWriter, ctx: Context[Params]) -> None:
    wanted = ctx.params.required
    if not raw.find("**/*.mcap"):
        ctx.not_applicable("NO_MCAP", "no .mcap file in this upload")
    topics = {t.name: t for t in raw.mcap().topics}
    missing = [name for name in wanted if name not in topics]
    out.result("topics_present", "fail" if missing else "pass", message=f"missing: {', '.join(missing)}" if missing else "all required topics found")
    slow = [name for name in wanted if name in topics and topics[name].rate < ctx.params.min_rate_hz]
    out.result("topic_rate", "warn" if slow else "pass", threshold=ctx.params.min_rate_hz, message=f"slow: {', '.join(slow)}" if slow else "rates are fine")

"""Check (episode level): is each episode long enough to train on, and does it say what the robot or person does?

In:  one recording and one episode of it (a start and an end on the recording's timeline, plus a task text).
Out: two check results for that episode, "pass", "warn" or "fail". The step itself still succeeds.
Why: very short episodes and episodes without a task text are the usual reasons a training run learns nothing,
     so this check flags them before the episode reaches a dataset version.
"""

from dataclasses import dataclass

from embodiflow import Check, CheckWriter, Context, Episode, Recording, algorithm, param


@dataclass
class Params:
    min_s: float = param(2.0, minimum=0.0, title="Shortest accepted episode (s)")


@algorithm.check(
    name="episode-duration",
    title="Episode duration and task",
    level="episode",
    checks=[
        Check("episode_duration", "The episode is long enough"),
        Check("episode_task", "The episode has a task text"),
    ],
    params=Params,
)
def run(recording: Recording, episode: Episode, out: CheckWriter, ctx: Context[Params]) -> None:
    # `episode.start_s` and `episode.end_s` are seconds on the recording's timeline.
    long_enough = episode.duration_s >= ctx.params.min_s
    out.result(
        "episode_duration",
        "pass" if long_enough else "fail",
        value=episode.duration_s,
        threshold=ctx.params.min_s,
        message=f"{episode.duration_s:.2f} s",
        at={"startS": episode.start_s, "endS": episode.end_s},
    )
    has_task = bool(episode.task and episode.task.strip())
    out.result("episode_task", "pass" if has_task else "warn", message=episode.task or "no task text")