captionizer.py

#
# SPDX-FileCopyrightText: 2026 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2026 This source file is part of the SensorTSLM open-source project.
#
# SPDX-License-Identifier: MIT
#
from __future__ import annotations

from timef.schema import CaptionResult
from transformer import Transformer
from annotator import Annotator
from reviewer import EvaluationResult, Reviewer


class Captionizer:
    def __init__(
        self,
        dataset,
        transformer: Transformer,
        annotator: Annotator,
        reviewer: Reviewer | None = None,
    ) -> None:
        self.dataset = dataset
        self.transformer = transformer
        self.annotator = annotator
        self.reviewer = reviewer

    def run(
        self, max_rows: int | None = None,
    ) -> tuple[CaptionResult, EvaluationResult | None]:
        result = CaptionResult()
        failed_rows: list[int] = []

        for i in range(len(self.dataset)):
            if max_rows is not None and i >= max_rows:
                break

            try:
                row = self.transformer.transform_row(self.dataset[i])
                row.annotations.extend(self.annotator.annotate(row))
                result.rows.append(row)
            except Exception:
                failed_rows.append(i)

        if failed_rows:
            print(f"{len(failed_rows)} / {i + 1} rows failed during captioning: {failed_rows}")

        evaluation = self.reviewer.evaluate(result) if self.reviewer else None
        return result, evaluation


if __name__ == "__main__":
    from mhc.dataset import MHCDataset
    from mhc.transformer import MHCTransformer
    from mhc.constants import MHC_CHANNEL_CONFIG
    from extractors.cross_channel import CrossChannelExtractor
    from synthesizers.sleep import SleepSynthesizer
    from extractors.statistical import StatisticalExtractor
    from extractors.structural import StructuralExtractor
    from models.local import LocalConfig, LocalModel
    from models.client import ClientConfig, ClientModel
    from extractors.semantic import SemanticExtractor
    from visualizer import plot_row
    import numpy as np

    dataset = MHCDataset(min_wear_pct=0.0)
    annotator = Annotator([
        StatisticalExtractor(MHC_CHANNEL_CONFIG),
        StructuralExtractor(MHC_CHANNEL_CONFIG),
        SemanticExtractor(MHC_CHANNEL_CONFIG),
        CrossChannelExtractor(MHC_CHANNEL_CONFIG, synthesizers=[SleepSynthesizer()]),
    ])

    captionizer = Captionizer(dataset, MHCTransformer(), annotator)
    print(f"Dataset size: {len(dataset)}")

    result, _ = captionizer.run(max_rows=1)
    print(f"Rows: {len(result.rows)}")
    print(f"Annotations: {sum(len(row.annotations) for row in result.rows)}")

    filtered = CaptionResult(rows=[row for row in result.iter_rows() if row.active_channel_count() >= 9])
    print(f"Filtered rows: {len(filtered.rows)}")

    model = ClientModel(
        ClientConfig(
            base_url="https://api.openai.com/v1",
            model="gpt-5.4",
        ),
        MHC_CHANNEL_CONFIG,
    )
    reviewer = Reviewer(model)
    evaluation = reviewer.evaluate(filtered)
    print(f"Evaluation: {len(evaluation.scores)} scores, mean={evaluation.mean_score}")
    for j, s in enumerate(evaluation.scores):
        print(f"  [{j}] score={s.score}  feedback={s.feedback}")

    for i, row in enumerate(filtered.iter_rows()):
        plot_row(row, save_path=f"sample_plot_{i}.png")
        print(f"Saved sample_plot_{i}.png")
        if i >= 3:
            break