标准化

"""Standardize: one MP4 and a per-frame CSV become the canonical recording.

In:  the uploaded files (raw): one .mp4 from the camera and one .csv with a header row of joint names and
     one row of joint values per video frame, recorded at 30 Hz.
Out: the canonical recording: one camera, one "observation.state" column, a task text.
Why: every later step (checks, curation, export, training) reads the canonical recording, never the raw files.
     Anything that does not fit (no video, two CSVs) is "not applicable", never guessed at.
"""

import csv
from dataclasses import dataclass

import numpy as np

from embodiflow import Context, RawFiles, RecordingBuilder, algorithm, param


@dataclass
class Params:
    task: str = param("move the arm", title="Task text", description="Recorded as the task of the episode")


@algorithm.standardize(
    name="video-csv",
    title="Video and CSV episode",
    description="Reads one .mp4 and one .csv with a header row and one numeric row per video frame, at 30 Hz.",
    params=Params,
)
def run(raw: RawFiles, out: RecordingBuilder, ctx: Context[Params]) -> None:
    videos = raw.find("**/*.mp4")
    tables = raw.find("**/*.csv")
    if len(videos) != 1 or len(tables) != 1:
        ctx.not_applicable("NOT_ONE_VIDEO_AND_CSV", f"expected one .mp4 and one .csv, found {len(videos)} and {len(tables)}")
    with open(tables[0].local_path, newline="", encoding="utf-8") as handle:
        reader = csv.reader(handle)
        names = next(reader)
        states = np.asarray([[float(cell) for cell in row] for row in reader], dtype=np.float32)
    out.fps(30)
    out.task(ctx.params.task)
    out.camera("observation.images.top", "top", video=videos[0].local_path)
    out.column("observation.state", states, "vector", names=names)
    ctx.progress(1, 1, "converted")

"""Standardize: one HDF5 file (frames, joint positions) becomes the canonical recording.

In:  the uploaded files (raw): one .h5 file with /frames (N,H,W,3 uint8) and /joint_positions (N,7).
Out: the canonical recording: one camera, one "observation.state" column in radians, a task text.
Why: teleoperation tools often dump everything into one HDF5 file; this shows how to read one with h5py
     and hand the arrays to the builder. Anything but exactly one .h5 file is "not applicable".
"""

from dataclasses import dataclass

import h5py
import numpy as np

from embodiflow import Context, RawFiles, RecordingBuilder, algorithm, param


@dataclass
class Params:
    task: str = param("move the arm", title="Task text", description="Used when the file carries no task attribute")


@algorithm.standardize(
    name="hdf5-arm",
    title="HDF5 arm episode",
    description="Reads /frames (N,H,W,3 uint8) and /joint_positions (N,7) from one .h5 file.",
    params=Params,
    requirements=["h5py>=3.10"],
)
def run(raw: RawFiles, out: RecordingBuilder, ctx: Context[Params]) -> None:
    files = raw.find("**/*.h5")
    if len(files) != 1:
        ctx.not_applicable("NOT_ONE_HDF5", f"expected one .h5 file, found {len(files)}")
    with h5py.File(files[0].local_path, "r") as h5:
        frames = np.asarray(h5["frames"])
        joints = np.asarray(h5["joint_positions"], dtype=np.float32)
        task = h5.attrs.get("task", ctx.params.task)
    out.fps(30)
    out.task(str(task))
    out.camera("observation.images.top", "top", frames=frames)
    out.column("observation.state", joints, "action.joint", unit="rad", frame="base", names=[f"j{i}" for i in range(joints.shape[1])])
    ctx.progress(1, 1, "converted")