diff --git a/src/jabs/scripts/cli/postprocessing.py b/src/jabs/scripts/cli/postprocessing.py index 1074776e..017b2020 100644 --- a/src/jabs/scripts/cli/postprocessing.py +++ b/src/jabs/scripts/cli/postprocessing.py @@ -27,7 +27,7 @@ import json import logging import shutil -from collections.abc import Callable +from collections.abc import Callable, Iterable from pathlib import Path from typing import Any @@ -287,6 +287,35 @@ def _list_behaviors(prediction_path: Path) -> list[str]: raise click.ClickException(f"Cannot open prediction file: {exc}") from exc +def _multiclass_behaviors(prediction_path: Path, behaviors: Iterable[str]) -> list[str]: + """Return which of the given behaviors hold multi-class predictions. + + A behavior group carries a ``class_names`` dataset only for multi-class + predictions, so its presence is what distinguishes them from binary ones. + + Args: + prediction_path: Path to the JABS prediction HDF5 file. + behaviors: Behavior names (safe-name HDF5 group keys) to check. + + Returns: + The subset of ``behaviors`` whose predictions are multi-class, in the + order given. + + Raises: + click.ClickException: If the file cannot be opened. + """ + try: + with h5py.File(prediction_path, "r") as h5: + prediction_group = h5["predictions"] + return [ + behavior + for behavior in behaviors + if behavior in prediction_group and "class_names" in prediction_group[behavior] + ] + except OSError as exc: + raise click.ClickException(f"Cannot open prediction file: {exc}") from exc + + def run_apply_postprocessing( prediction_file: Path, config: list[dict[str, Any]] | _BehaviorConfigMap, @@ -351,6 +380,19 @@ def run_apply_postprocessing( f"Available behaviors: {', '.join(repr(b) for b in available_in_file)}" ) + # The stages classify each frame as behavior or not behavior (see + # ClassLabels), so they have no meaning for multi-class predictions. The GUI + # gates on this too: MultiClassClassifyStrategy.postprocess_identity returns + # None. Refuse here, before the output file is created, rather than writing + # results that treat one class as "the behavior" and the rest as background. + multiclass = _multiclass_behaviors(prediction_file, behavior_configs) + if multiclass: + raise click.ClickException( + f"Post-processing supports binary predictions only, but the following " + f"behavior(s) hold multi-class predictions: " + f"{', '.join(repr(b) for b in multiclass)}." + ) + # --- Prepare output file ------------------------------------------------- if output_path != prediction_file: try: @@ -397,6 +439,9 @@ def run_apply_postprocessing( predicted_class_postprocessed=postprocessed, identity_to_track=pred.identity_to_track, external_identity_mapping=pred.external_identity_mapping, + # carried through so writing the result cannot drop it: the HDF5 + # adapter deletes an existing class_names dataset when this is None + class_names=pred.class_names, ) try: io.save(updated_pred, output_path) diff --git a/tests/scripts/test_apply_postprocessing.py b/tests/scripts/test_apply_postprocessing.py index 487254f2..fe7ce08e 100644 --- a/tests/scripts/test_apply_postprocessing.py +++ b/tests/scripts/test_apply_postprocessing.py @@ -379,3 +379,105 @@ def test_run_preserves_raw_predictions(tmp_path: Path) -> None: updated: BehaviorPrediction = io.load(pred_file, BehaviorPrediction, behavior="grooming") np.testing.assert_array_equal(updated.predicted_class, original.predicted_class) np.testing.assert_array_equal(updated.probabilities, original.probabilities) + + +# --------------------------------------------------------------------------- +# multi-class predictions +# --------------------------------------------------------------------------- + + +def _make_multiclass_prediction( + behavior: str, + class_names: list[str], + n_identities: int = _N_IDENTITIES, + n_frames: int = _N_FRAMES, +) -> BehaviorPrediction: + """Build a multi-class BehaviorPrediction (3-D probabilities, class names set).""" + rng = np.random.default_rng(seed=7) + return BehaviorPrediction( + behavior=behavior, + predicted_class=rng.integers( + 0, len(class_names), size=(n_identities, n_frames), dtype=np.int8 + ), + probabilities=rng.random((n_identities, n_frames, len(class_names))).astype(np.float32), + classifier=_CLASSIFIER, + pose_file="pose.h5", + pose_hash="deadbeef", + class_names=class_names, + ) + + +def test_run_multiclass_predictions_raises(tmp_path: Path) -> None: + """A multi-class prediction file is refused with an explanatory message. + + The stages label frames behavior/not-behavior, so they have no meaning for + multi-class predictions. Previously this reached the write step and died with an + opaque "probabilities shape ... does not match expected shape" ValueError. + """ + import click + + pred_file = tmp_path / "predictions.h5" + io.save(_make_multiclass_prediction("multiclass", ["None", "walking", "grooming"]), pred_file) + + config = [{"stage_name": "BoutDurationFilterStage", "parameters": {"min_duration": 3}}] + + with pytest.raises(click.ClickException, match="binary predictions only"): + run_apply_postprocessing(pred_file, config, "multiclass", pred_file) + + +def test_run_multiclass_refusal_leaves_the_file_untouched(tmp_path: Path) -> None: + """The refusal happens before any output is written.""" + import click + + pred_file = tmp_path / "predictions.h5" + out_file = tmp_path / "out.h5" + io.save(_make_multiclass_prediction("multiclass", ["None", "walking"]), pred_file) + before = pred_file.read_bytes() + + config = [{"stage_name": "BoutDurationFilterStage", "parameters": {"min_duration": 3}}] + + with pytest.raises(click.ClickException, match="binary predictions only"): + run_apply_postprocessing(pred_file, config, "multiclass", out_file) + + assert not out_file.exists() + assert pred_file.read_bytes() == before + + +def test_run_multiclass_behavior_named_alongside_binary_ones_raises(tmp_path: Path) -> None: + """A multi-class behavior in a multi-behavior config stops the whole run. + + Failing before the output file is created keeps a partially processed result + from being written. + """ + import click + + pred_file = tmp_path / "predictions.h5" + out_file = tmp_path / "out.h5" + _write_prediction_file(pred_file, "grooming") + io.save(_make_multiclass_prediction("multiclass", ["None", "walking"]), pred_file) + + stages = [{"stage_name": "BoutDurationFilterStage", "parameters": {"min_duration": 3}}] + config = {"grooming": stages, "multiclass": stages} + + with pytest.raises(click.ClickException, match="multiclass"): + run_apply_postprocessing(pred_file, config, None, out_file) + + assert not out_file.exists() + + +def test_run_binary_behavior_in_a_file_that_also_has_multiclass(tmp_path: Path) -> None: + """Only the named behaviors are checked, so a binary one still processes.""" + pred_file = tmp_path / "predictions.h5" + out_file = tmp_path / "out.h5" + _write_prediction_file(pred_file, "grooming") + io.save(_make_multiclass_prediction("multiclass", ["None", "walking"]), pred_file) + + config = [{"stage_name": "BoutDurationFilterStage", "parameters": {"min_duration": 3}}] + processed = run_apply_postprocessing(pred_file, config, "grooming", out_file) + + assert processed == ["grooming"] + updated: BehaviorPrediction = io.load(out_file, BehaviorPrediction, behavior="grooming") + assert updated.predicted_class_postprocessed is not None + # the untouched multi-class group keeps its class names + multiclass: BehaviorPrediction = io.load(out_file, BehaviorPrediction, behavior="multiclass") + assert multiclass.class_names == ["None", "walking"]