Task/normalize behavior name resolution

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jabs-postprocess"
-version = "0.5.3"
+version = "0.5.4"
 description = "A python library for JABS postprocessing utilities."
 readme = "README.md"
 license = "LicenseRef-PLATFORM-LICENSE-AGREEMENT-FOR-NON-COMMERCIAL-USE"

src/jabs_postprocess/utils/project_utils.py

@@ -28,6 +28,28 @@
 BEHAVIOR_CLASSIFY_VERSION = 1
 
 
+def normalize_behavior_name(behavior: str) -> str:
+    """Normalize behavior names to align with pipeline behavior path rules."""
+    normalized = behavior.replace(" ", "_")
+    normalized = re.sub(r"[()]", "", normalized)
+    return normalized.lower()
+
+
+def resolve_behavior_key(behavior: str, available_behaviors: List[str]) -> str | None:
+    """Resolves a behavior key from available behaviors using normalization."""
+    if behavior in available_behaviors:
+        return behavior
+    behavior_norm = normalize_behavior_name(behavior)
+    matches = [
+        key
+        for key in available_behaviors
+        if normalize_behavior_name(key) == behavior_norm
+    ]
+    if len(matches) == 1:
+        return matches[0]
+    return None
+
+
 class MissingBehaviorException(ValueError):
     """Custom error for behavior-related missing data."""
 
@@ -705,10 +727,11 @@ def from_jabs_annotation_file(cls, source_file: Path, behavior: str):
             data = json.load(f)
 
         vid_name = data["file"]
+        behavior_norm = normalize_behavior_name(behavior)
         df_list = []
         for animal_idx, labels in data["labels"].items():
             for cur_behavior, label_data in labels.items():
-                if cur_behavior == behavior:
+                if normalize_behavior_name(cur_behavior) == behavior_norm:
                     new_events = []
                     for cur_event in label_data:
                         new_df = pd.DataFrame(
@@ -1340,8 +1363,9 @@ def generate_bout_table(source_file: Path, settings: ClassifierSettings):
                 MissingBehaviorException if behavior file exists but contains no behavior predictions.
         """
         with h5py.File(str(source_file), "r") as in_f:
-            if settings.behavior not in in_f["predictions/"].keys():
-                available_keys = list(in_f["predictions"].keys())
+            available_keys = list(in_f["predictions"].keys())
+            behavior_key = resolve_behavior_key(settings.behavior, available_keys)
+            if behavior_key is None:
                 if len(available_keys) > 0:
                     behavior_pred_shape = in_f[
                         f"predictions/{available_keys[0]}/predicted_class"
@@ -1353,7 +1377,7 @@ def generate_bout_table(source_file: Path, settings: ClassifierSettings):
                     raise MissingBehaviorException(
                         "Prediction file exists, but no behaviors present to discover shape."
                     )
-            class_calls = in_f[f"predictions/{settings.behavior}/predicted_class"][:]
+            class_calls = in_f[f"predictions/{behavior_key}/predicted_class"][:]
 
         # Iterate over the animals
         bout_dfs = []

tests/utils/test_behavior_normalization.py

@@ -0,0 +1,78 @@
+"""Tests for behavior name normalization and resolution."""
+
+import json
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+from jabs_postprocess.utils.project_utils import (
+    BoutTable,
+    ClassifierSettings,
+    Prediction,
+    normalize_behavior_name,
+    resolve_behavior_key,
+)
+
+
+def test_normalize_behavior_name_matches_pipeline_rules():
+    assert normalize_behavior_name("rearing (unsupported)") == "rearing_unsupported"
+    assert normalize_behavior_name("Rearing_unsupported") == "rearing_unsupported"
+
+
+def test_resolve_behavior_key_normalized_match():
+    available = ["rearing (unsupported)", "grooming"]
+    assert (
+        resolve_behavior_key("Rearing_unsupported", available)
+        == "rearing (unsupported)"
+    )
+
+
+def test_resolve_behavior_key_ambiguous_match():
+    available = ["A B", "A_B"]
+    assert resolve_behavior_key("a b", available) is None
+
+
+def test_from_jabs_annotation_file_uses_normalized_behavior(tmp_path: Path):
+    annotation = {
+        "file": "video1.avi",
+        "labels": {
+            "0": {
+                "rearing (unsupported)": [
+                    {"start": 0, "end": 2, "present": 1}
+                ]
+            }
+        },
+    }
+    annotation_path = tmp_path / "annotation.json"
+    annotation_path.write_text(json.dumps(annotation))
+
+    table = BoutTable.from_jabs_annotation_file(
+        annotation_path, "Rearing_unsupported"
+    )
+    df = table.data
+
+    assert len(df) == 1
+    row = df.iloc[0]
+    assert row["start"] == 0
+    assert row["duration"] == 3
+    assert row["is_behavior"] == 1
+    assert row["video_name"] == "video1"
+
+
+def test_generate_bout_table_uses_normalized_behavior(tmp_path: Path):
+    pred_path = tmp_path / "predictions.h5"
+    with h5py.File(pred_path, "w") as h5f:
+        pred_group = h5f.create_group("predictions")
+        behavior_group = pred_group.create_group("rearing (unsupported)")
+        behavior_group.create_dataset(
+            "predicted_class", data=np.array([[0, 1, 1, 0]], dtype=np.int8)
+        )
+
+    settings = ClassifierSettings("Rearing_unsupported", 0, 0, 0)
+    df = Prediction.generate_bout_table(pred_path, settings)
+
+    match = df[
+        (df["start"] == 1) & (df["duration"] == 2) & (df["is_behavior"] == 1)
+    ]
+    assert len(match) == 1