增强

"""Enrich: the mean absolute difference between consecutive frames of the first camera, per frame.

In:  the canonical recording.
Out: one new per-frame feature "motion_energy"; the recording itself is not changed.
Why: idle stretches (the arm waits, the person stands still) show as low motion, so a later step can
     segment or drop them. A feature is how an algorithm adds knowledge without rewriting the recording.
"""

from dataclasses import dataclass

import numpy as np

from embodiflow import Context, Feature, FeatureWriter, Recording, algorithm, param


@dataclass
class Params:
    step: int = param(1, minimum=1, title="Frame step")


@algorithm.enrich(
    name="motion-energy",
    title="Motion energy",
    description="How much the first camera's picture changes from one frame to the next.",
    features=[Feature("motion_energy", type="scalar", per="frame", dtype="float32", unit="none")],
    params=Params,
)
def run(recording: Recording, out: FeatureWriter, ctx: Context[Params]) -> None:
    cam = recording.cameras[0]
    energy = np.zeros((recording.num_frames, 1), np.float32)
    prev = None
    for frame in cam.frames(step=ctx.params.step, gray=True):
        current = frame.image.astype(np.float32)
        if prev is not None:
            energy[frame.index, 0] = np.abs(current - prev).mean()
        prev = current
    out.per_frame("motion_energy", energy)

"""Enrich (episode level): the length of each episode and whether it is long enough to train on.

In:  the canonical recording and one episode of it, the step runs once per episode.
Out: two per-episode features: "length_s" and "is_long".
Why: shows the episode level: the step gets the episode window, not the whole recording.
"""

from dataclasses import dataclass

from embodiflow import Context, Episode, Feature, FeatureWriter, Recording, algorithm, param


@dataclass
class Params:
    long_s: float = param(1.0, minimum=0.0, title="Seconds an episode needs to count as long")


@algorithm.enrich(
    name="episode-length",
    title="Episode length",
    level="episode",
    features=[
        Feature("length_s", type="scalar", per="episode", dtype="float64", unit="s"),
        Feature("is_long", type="label.bool", per="episode"),
    ],
    params=Params,
)
def run(recording: Recording, episode: Episode, out: FeatureWriter, ctx: Context[Params]) -> None:
    length = episode.duration_s  # seconds, on the recording's timeline
    out.per_episode("length_s", length)
    out.per_episode("is_long", length >= ctx.params.long_s)