feat(tlabel): add TLabel format reader/writer

forge/formats/__init__.py

@@ -56,6 +56,12 @@ def _register_formats() -> None:
     except ImportError:
         pass
 
+    try:
+        from forge.formats import tlabel  # noqa: F401
+    except ImportError:
+        pass  # TLabel dependencies not installed
+
 
 # Register formats on module import
 _register_formats()
+

forge/formats/tlabel/__init__.py

@@ -0,0 +1,10 @@
+"""TLabel format support for Forge.
+
+Provides reading and writing of .tlabel tactile annotation files.
+TLabel (https://github.com/liesliy/tlabel) is an open standard for
+tactile sensor data with a 14-dimensional semantic schema.
+"""
+
+from forge.formats.tlabel.reader import TLabelReader
+
+__all__ = ["TLabelReader"]

forge/formats/tlabel/reader.py

@@ -0,0 +1,221 @@
+"""TLabel format reader for Forge.
+
+Reads .tlabel JSON files containing tactile sensor annotations.
+TLabel Schema V2 defines 14 semantic dimensions for tactile data
+(contact, force, slip, vibration, texture, etc.).
+
+See: https://github.com/liesliy/tlabel
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from forge.core.exceptions import (
+    EpisodeNotFoundError,
+    InspectionError,
+)
+from forge.core.models import (
+    DatasetInfo,
+    Dtype,
+    Episode,
+    FieldSchema,
+    Frame,
+)
+from forge.formats.registry import FormatRegistry
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
+def _load_tlabel(path: Path) -> dict[str, Any]:
+    """Load and validate a .tlabel JSON file."""
+    if not path.exists():
+        raise FileNotFoundError(f"TLabel file not found: {path}")
+
+    with open(path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    # Basic structure check
+    if "frames" not in data and "frame" not in data:
+        raise InspectionError(f"Not a valid TLabel file: missing 'frames' field")
+
+    return data
+
+
+def _frame_to_forge(tlabel_frame: dict[str, Any]) -> Frame:
+    """Convert a single TLabel frame dict to a Forge Frame."""
+    schema_v2 = tlabel_frame.get("schema_v2", {})
+
+    return Frame(
+        index=tlabel_frame.get("frame_idx", 0),
+        timestamp=tlabel_frame.get("timestamp_s"),
+        is_first=tlabel_frame.get("is_first", False),
+        is_last=tlabel_frame.get("is_last", False),
+        extras={
+            "tlabel.schema_v2": schema_v2,
+            "tlabel.confidence": tlabel_frame.get("confidence"),
+            "tlabel.manipulation_phase": tlabel_frame.get("manipulation_phase"),
+        },
+    )
+
+
+@FormatRegistry.register_reader("tlabel")
+class TLabelReader:
+    """Reader for .tlabel tactile annotation files.
+
+    TLabel files contain per-frame tactile sensor annotations with a
+    14-dimensional semantic schema (Schema V2). Each frame maps to a
+    Forge Frame with tactile data stored in Frame.extras.
+    """
+
+    @property
+    def format_name(self) -> str:
+        return "tlabel"
+
+    @classmethod
+    def can_read(cls, path: Path) -> bool:
+        """Check if path looks like a .tlabel file or directory containing them."""
+        path = Path(path)
+
+        # Single .tlabel file
+        if path.is_file() and path.suffix == ".tlabel":
+            return True
+
+        # Directory with .tlabel files
+        if path.is_dir():
+            return any(path.glob("*.tlabel"))
+
+        # JSON file with TLabel structure (some use .json extension)
+        if path.is_file() and path.suffix == ".json":
+            try:
+                with open(path, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                return "frames" in data and "schema_version" in data
+            except (json.JSONDecodeError, OSError):
+                return False
+
+        return False
+
+    @classmethod
+    def detect_version(cls, path: Path) -> str | None:
+        """Detect TLabel schema version."""
+        path = Path(path)
+        if path.is_file():
+            try:
+                with open(path, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                return data.get("schema_version_v2", data.get("schema_version"))
+            except (json.JSONDecodeError, OSError):
+                return None
+        return None
+
+    def inspect(self, path: Path) -> DatasetInfo:
+        """Analyze a .tlabel file's structure."""
+        path = Path(path)
+        data = _load_tlabel(path)
+
+        frames = data.get("frames", [])
+        num_frames = len(frames)
+        sensor = data.get("sensor", {})
+
+        # Build observation schema from capabilities
+        capabilities = data.get("capabilities", {})
+        obs_schema = {}
+        for dim_name, enabled in capabilities.items():
+            if enabled:
+                obs_schema[f"tlabel.{dim_name}"] = FieldSchema(
+                    name=f"tlabel.{dim_name}",
+                    shape=(1,),
+                    dtype=Dtype.FLOAT32,
+                    description=f"TLabel dimension: {dim_name}",
+                )
+
+        # Infer FPS from timestamps
+        inferred_fps = None
+        if num_frames >= 2:
+            t0 = frames[0].get("timestamp_s", 0)
+            t1 = frames[1].get("timestamp_s", 0)
+            dt = t1 - t0
+            if dt > 0:
+                inferred_fps = round(1.0 / dt, 1)
+
+        return DatasetInfo(
+            path=path,
+            format="tlabel",
+            format_version=data.get("schema_version"),
+            num_episodes=1,  # One file = one episode
+            total_frames=num_frames,
+            observation_schema=obs_schema,
+            cameras={},
+            has_timestamps=all("timestamp_s" in f for f in frames[:10]),
+            has_language=data.get("episode", {}).get("task") is not None,
+            has_rewards=False,
+            has_success_labels=data.get("episode", {}).get("success") is not None,
+            inferred_fps=inferred_fps,
+            inferred_robot_type=sensor.get("name"),
+            sample_num_frames=num_frames,
+            sample_language=data.get("episode", {}).get("task"),
+        )
+
+    def read_episodes(self, path: Path) -> Iterator[Episode]:
+        """Yield episodes from a .tlabel file (one file = one episode)."""
+        path = Path(path)
+
+        # If directory, iterate .tlabel files
+        if path.is_dir():
+            for tlabel_file in sorted(path.glob("*.tlabel")):
+                yield from self._read_single(tlabel_file)
+            for json_file in sorted(path.glob("*.json")):
+                if self.can_read(json_file):
+                    yield from self._read_single(json_file)
+        else:
+            yield from self._read_single(path)
+
+    def read_episode(self, path: Path, episode_id: str) -> Episode:
+        """Read a specific episode by ID."""
+        for ep in self.read_episodes(path):
+            if ep.episode_id == episode_id:
+                return ep
+        raise EpisodeNotFoundError(episode_id, str(path))
+
+    def _read_single(self, file_path: Path) -> Iterator[Episode]:
+        """Read a single .tlabel file as one Episode."""
+        data = _load_tlabel(file_path)
+        episode_meta = data.get("episode", {})
+        frames_data = data.get("frames", [])
+
+        # Build Episode
+        ep_id = episode_meta.get("id", file_path.stem)
+        episode = Episode(
+            episode_id=ep_id,
+            metadata={
+                "sensor": data.get("sensor", {}),
+                "capabilities": data.get("capabilities", {}),
+                "schema_version": data.get("schema_version"),
+            },
+            language_instruction=episode_meta.get("task"),
+            success=episode_meta.get("success"),
+            fps=self._infer_fps(frames_data),
+        )
+
+        # Lazy frame loader
+        def _load_frames() -> Iterator[Frame]:
+            for fd in frames_data:
+                yield _frame_to_forge(fd)
+
+        episode._frame_loader = _load_frames
+        yield episode
+
+    @staticmethod
+    def _infer_fps(frames_data: list[dict]) -> float | None:
+        if len(frames_data) >= 2:
+            t0 = frames_data[0].get("timestamp_s", 0)
+            t1 = frames_data[1].get("timestamp_s", 0)
+            dt = t1 - t0
+            if dt > 0:
+                return round(1.0 / dt, 1)
+        return None

forge/formats/tlabel/writer.py

@@ -0,0 +1,108 @@
+"""TLabel format writer for Forge.
+
+Writes Forge Episodes to .tlabel JSON files with Schema V2 annotations.
+Tactile data stored in Frame.extras["tlabel.schema_v2"] is preserved.
+
+See: https://github.com/liesliy/tlabel
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from forge.core.models import DatasetInfo, Episode
+from forge.formats.registry import FormatRegistry
+
+if TYPE_CHECKING:
+    pass
+
+
+def _frame_to_tlabel(frame_obj: Any) -> dict[str, Any]:
+    """Convert a Forge Frame to a TLabel frame dict."""
+    extras = getattr(frame_obj, "extras", {}) or {}
+    schema_v2 = extras.get("tlabel.schema_v2", {})
+
+    return {
+        "frame_idx": frame_obj.index,
+        "timestamp_s": frame_obj.timestamp,
+        "schema_v2": schema_v2,
+        "confidence": extras.get("tlabel.confidence"),
+        "manipulation_phase": extras.get("tlabel.manipulation_phase"),
+        "is_first": getattr(frame_obj, "is_first", False),
+        "is_last": getattr(frame_obj, "is_last", False),
+    }
+
+
+@FormatRegistry.register_writer("tlabel")
+class TLabelWriter:
+    """Writer for .tlabel tactile annotation files.
+
+    Converts Forge Episodes back to .tlabel JSON format.
+    Tactile data from Frame.extras["tlabel.schema_v2"] is written
+    to the standard TLabel Schema V2 frame structure.
+    """
+
+    @property
+    def format_name(self) -> str:
+        return "tlabel"
+
+    def write_episode(
+        self,
+        episode: Episode,
+        output_path: Path,
+        episode_index: int | None = None,
+    ) -> None:
+        """Write a single episode to a .tlabel file."""
+        output_path = Path(output_path)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        ep_suffix = f"_{episode_index:04d}" if episode_index is not None else ""
+        out_file = output_path / f"{episode.episode_id}{ep_suffix}.tlabel"
+
+        frames = [_frame_to_tlabel(f) for f in episode.frames()]
+
+        data = {
+            "schema_version": episode.metadata.get("schema_version", "0.17.0"),
+            "sensor": episode.metadata.get("sensor", {}),
+            "capabilities": episode.metadata.get("capabilities", {}),
+            "episode": {
+                "id": episode.episode_id,
+                "task": episode.language_instruction,
+                "num_frames": len(frames),
+                "success": episode.success,
+            },
+            "frames": frames,
+        }
+
+        with open(out_file, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, ensure_ascii=False, default=str)
+
+    def write_dataset(
+        self,
+        episodes: Iterator[Episode],
+        output_path: Path,
+        dataset_info: DatasetInfo | None = None,
+    ) -> None:
+        """Write multiple episodes to .tlabel files."""
+        output_path = Path(output_path)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        for i, episode in enumerate(episodes):
+            self.write_episode(episode, output_path, episode_index=i)
+
+    def finalize(self, output_path: Path, dataset_info: DatasetInfo) -> None:
+        """Write a summary manifest.json for the dataset."""
+        output_path = Path(output_path)
+        tlabel_files = sorted(output_path.glob("*.tlabel"))
+
+        manifest = {
+            "format": "tlabel",
+            "num_episodes": len(tlabel_files),
+            "files": [f.name for f in tlabel_files],
+        }
+
+        with open(output_path / "manifest.json", "w", encoding="utf-8") as f:
+            json.dump(manifest, f, indent=2)