Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "jabs-postprocess"
version = "0.3.1"
version = "0.4.0"
description = "A python library for JABS postprocessing utilities."
readme = "README.md"
license = "LicenseRef-PLATFORM-LICENSE-AGREEMENT-FOR-NON-COMMERCIAL-USE"
Expand Down
99 changes: 66 additions & 33 deletions src/jabs_postprocess/cli/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
DEFAULT_MIN_BOUT,
DEFAULT_STITCH,
)
from jabs_postprocess.cli.utils import load_json

app = typer.Typer()

Expand Down Expand Up @@ -199,9 +200,6 @@ def generate_tables(
help="Folder that contains the project with both pose files and behavior prediction files"
),
],
behavior: Annotated[
List[str], typer.Option(help="Behavior(s) to produce table(s) for")
],
out_prefix: Annotated[
str,
typer.Option(
Expand All @@ -217,37 +215,38 @@ def generate_tables(
help="If features were exported, include feature-based characteristics of bouts"
),
] = None,
interpolate_size: Annotated[
Optional[int],
typer.Option(
help=f"Maximum number of frames in which missing data will be interpolated (default: {DEFAULT_INTERPOLATE})"
),
] = None,
stitch_gap: Annotated[
Optional[int],
typer.Option(
help=f"Number of frames in which sequential behavior prediction bouts will be joined (default: {DEFAULT_STITCH})"
),
] = None,
min_bout_length: Annotated[
Optional[int],
typer.Option(
help=f"Minimum number of frames in which a behavior prediction must be to be considered (default: {DEFAULT_MIN_BOUT})"
),
] = None,
overwrite: Annotated[bool, typer.Option(help="Overwrites output files")] = False,
add_statistics: Annotated[
bool,
typer.Option(
help="Add bout statistics (count, duration stats, latency) to behavior tables",
),
] = True,
behavior_config: Path | None = typer.Option(
None, "--behavior-config", help="JSON file with behavior configurations"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This could benefit from indicating some structure of the JSON, e.g. what fields it expects and which are optional.

),
behaviors: List[str] | None = typer.Option(
None, "--behavior", help="Simple behavior names (uses defaults)"
),
):
"""Generate behavior tables from JABS predictions.

This command transforms behavior predictions from a JABS project into tabular format,
creating both bout-level and summary tables.

Example JSON for behavior_config argument:
{
"behaviors": [
{"behavior": "Behavior_1_Name", "interpolate_size": 1},
{"behavior": "Behavior_2_Name", "stitch_gap": 30, "min_bout_length": 150}
{"behavior": "Behavior_2_Name",
"stitch_gap": 30,
"min_bout_length": 150,
"interpolate_size": 1
}
]
}

The --add-statistics option adds additional columns with bout-level statistics:
- total_bout_count: Number of behavior bouts per animal
- avg_bout_duration: Average bout duration per animal
Expand All @@ -258,30 +257,62 @@ def generate_tables(
# Convert Path to string
feature_folder = feature_folder if feature_folder else None

behaviors = []
for behavior_name in behavior:
behavior_config = {
"behavior": behavior_name,
"interpolate_size": interpolate_size,
"stitch_gap": stitch_gap,
"min_bout_length": min_bout_length,
}
behaviors.append(behavior_config)
behavior_args = []

if behavior_config:
try:
config = load_json(behavior_config)
if not isinstance(config, dict) or "behaviors" not in config:

@gbeane gbeane Oct 8, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you could use jsonschema.validate() to validate the json (I don't think this needs to change for this pull request).

raise ValueError("Config must be a JSON object with 'behaviors' key")

for b in config["behaviors"]:
behavior_args.append(
{
"behavior": b["behavior"],
"interpolate_size": b.get(
"interpolate_size", DEFAULT_INTERPOLATE
),
"stitch_gap": b.get("stitch_gap", DEFAULT_STITCH),
"min_bout_length": b.get("min_bout_length", DEFAULT_MIN_BOUT),
}
)
except (ValueError, TypeError, KeyError) as e:
typer.echo(f"Error loading behavior config: {e}", err=True)
raise typer.Exit(1)
elif behaviors:
for behavior_name in behaviors:
behavior_args.append(
{
"behavior": behavior_name,
"interpolate_size": DEFAULT_INTERPOLATE,
"stitch_gap": DEFAULT_STITCH,
"min_bout_length": DEFAULT_MIN_BOUT,
}
)
else:
typer.echo(
"Error: Must provide either --behavior-config or --behavior options",
err=True,
)
raise typer.Exit(1)

results = generate_behavior_tables.process_multiple_behaviors(
project_folder=project_folder,
behaviors=behaviors,
behaviors=behavior_args,
out_prefix=out_prefix,
out_bin_size=out_bin_size,
feature_folder=feature_folder,
overwrite=overwrite,
)

# Extract behavior names for output
behavior_names = [b["behavior"] for b in behavior_args]

# Add bout statistics if requested
if add_statistics:
typer.echo("Adding bout statistics to generated tables...")
for behavior_name, (bout_file, summary_file) in zip(
behavior, results, strict=True
behavior_names, results, strict=True
):
try:
# Load bout table and add statistics
Expand All @@ -294,7 +325,9 @@ def generate_tables(
f" Warning: Failed to add statistics to {bout_file}: {str(e)}"
)

for behavior_name, (bout_file, summary_file) in zip(behavior, results, strict=True):
for behavior_name, (bout_file, summary_file) in zip(
behavior_names, results, strict=True
):
typer.echo(f"Generated tables for {behavior_name}:")
typer.echo(f" Bout table: {bout_file}")
typer.echo(f" Summary table: {summary_file}")
Expand Down
75 changes: 75 additions & 0 deletions src/jabs_postprocess/cli/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import json
from pathlib import Path
from typing import Dict, Any, List


def load_json(source: Path | str | None) -> Dict[str, Any] | List[Any] | None:
"""Load JSON from a file path, JSON string, or return None.

Args:
source: Can be:
- Path object pointing to a JSON file
- String containing a file path to a JSON file
- String containing raw JSON content
- None (returns None)

Returns:
Parsed JSON (dict, list, or other JSON-serializable type), or None if input is None

Raises:
ValueError: If source cannot be parsed as JSON or loaded from file
TypeError: If source is not a supported type
"""
if source is None:
return None

if isinstance(source, Path):
# It's a Path object, read from file
try:
with open(source, "r", encoding="utf-8") as f:
return json.load(f)
except FileNotFoundError:
raise ValueError(f"File not found: {source}")
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON in file {source}: {e}")

elif isinstance(source, str):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe the isinstance(source, str) block could be simplified by checking Path(source).exists() first to see if it is a path, then fall back to parsing the string as JSON and if that fails finally check if it is one of the quoted strings?

I think this could avoid the multiple attempts to open it as a file

# It's a string - could be a file path or JSON content
stripped = source.strip()

# Check if it looks like JSON (starts with {, [, or is a quoted string)
if stripped and stripped[0] in '{["' or stripped in ("true", "false", "null"):
# Try as JSON first
try:
return json.loads(source)
except json.JSONDecodeError:
# If it fails, try as file path
try:
with open(source, "r", encoding="utf-8") as f:
return json.load(f)
except FileNotFoundError:
# Original JSON parse error is more relevant
raise ValueError(
f"Invalid JSON: {source[:100]}{'...' if len(source) > 100 else ''}"
)
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON in file {source}: {e}")
else:
# Doesn't look like JSON, try as file path first
try:
with open(source, "r", encoding="utf-8") as f:
return json.load(f)
except FileNotFoundError:
# Last attempt: maybe it's JSON without typical start chars
try:
return json.loads(source)
except json.JSONDecodeError:
raise ValueError(
f"'{source}' is neither a valid file path nor valid JSON"
)
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON in file {source}: {e}")
else:
raise TypeError(
f"source must be Path, str, or None, got {type(source).__name__}"
)
76 changes: 57 additions & 19 deletions tests/cli/test_generate_tables.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,18 @@
3. Parameter validation and error handling
4. File output and overwrite behavior
5. Integration with underlying processing modules
6. JSON config file loading for behavior parameters
"""

from unittest.mock import MagicMock, patch
import pytest

from jabs_postprocess.cli.main import app
from jabs_postprocess.utils.metadata import (
DEFAULT_INTERPOLATE,
DEFAULT_MIN_BOUT,
DEFAULT_STITCH,
)


class TestGenerateTables:
Expand Down Expand Up @@ -89,11 +95,12 @@ def test_generate_tables_basic(
assert call_args.kwargs["out_prefix"] == out_prefix
assert len(call_args.kwargs["behaviors"]) == behavior_count

# When using --behavior (simple names), default values should be used
for i, behavior_config in enumerate(call_args.kwargs["behaviors"]):
assert behavior_config["behavior"] == f"behavior{i}"
assert "interpolate_size" in behavior_config
assert "stitch_gap" in behavior_config
assert "min_bout_length" in behavior_config
assert behavior_config["interpolate_size"] == DEFAULT_INTERPOLATE
assert behavior_config["stitch_gap"] == DEFAULT_STITCH
assert behavior_config["min_bout_length"] == DEFAULT_MIN_BOUT

# Verify statistics handling
if add_statistics or add_statistics is None:
Expand All @@ -112,8 +119,10 @@ def test_generate_tables_basic(
],
)
@patch("jabs_postprocess.cli.main.generate_behavior_tables")
def test_generate_tables_with_parameters(
@patch("jabs_postprocess.cli.main.load_json")
def test_generate_tables_with_json_config(
self,
mock_load_json,
mock_generate_module,
runner,
mock_project_folder,
Expand All @@ -122,10 +131,12 @@ def test_generate_tables_with_parameters(
stitch_gap,
min_bout_length,
out_bin_size,
tmp_path,
):
"""Test table generation with various parameter combinations.
"""Test table generation with JSON config file for custom parameters.

Args:
mock_load_json: Mock load_json function
mock_generate_module: Mock generate_behavior_tables module
runner: CLI test runner
mock_project_folder: Mock project directory
Expand All @@ -134,32 +145,47 @@ def test_generate_tables_with_parameters(
stitch_gap: Stitch gap parameter
min_bout_length: Minimum bout length parameter
out_bin_size: Output bin size parameter
tmp_path: Temporary directory for config file
"""
# Arrange
behavior = "test_behavior"
mock_generate_module.process_multiple_behaviors.return_value = [
("bout.csv", "summary.csv")
]

# Create behavior config with optional parameters
behavior_config_data = {
"behaviors": [
{
"behavior": behavior,
}
]
}
if interpolate_size is not None:
behavior_config_data["behaviors"][0]["interpolate_size"] = interpolate_size
if stitch_gap is not None:
behavior_config_data["behaviors"][0]["stitch_gap"] = stitch_gap
if min_bout_length is not None:
behavior_config_data["behaviors"][0]["min_bout_length"] = min_bout_length

# Mock the load_json function
mock_load_json.return_value = behavior_config_data

# Create config file path
config_file = tmp_path / "behavior_config.json"

cmd_args = [
"generate-tables",
"--project-folder",
str(mock_project_folder),
"--behavior",
behavior,
"--behavior-config",
str(config_file),
"--feature-folder",
str(mock_feature_folder),
"--out-bin-size",
str(out_bin_size),
]

if interpolate_size is not None:
cmd_args.extend(["--interpolate-size", str(interpolate_size)])
if stitch_gap is not None:
cmd_args.extend(["--stitch-gap", str(stitch_gap)])
if min_bout_length is not None:
cmd_args.extend(["--min-bout-length", str(min_bout_length)])

# Act
result = runner.invoke(app, cmd_args)

Expand All @@ -171,9 +197,18 @@ def test_generate_tables_with_parameters(
assert call_args.kwargs["out_bin_size"] == out_bin_size

behavior_config = call_args.kwargs["behaviors"][0]
assert behavior_config["interpolate_size"] == interpolate_size
assert behavior_config["stitch_gap"] == stitch_gap
assert behavior_config["min_bout_length"] == min_bout_length
# Check that parameters were set correctly (or use defaults if None)
expected_interpolate = (
interpolate_size if interpolate_size is not None else DEFAULT_INTERPOLATE
)
expected_stitch = stitch_gap if stitch_gap is not None else DEFAULT_STITCH
expected_min_bout = (
min_bout_length if min_bout_length is not None else DEFAULT_MIN_BOUT
)

assert behavior_config["interpolate_size"] == expected_interpolate
assert behavior_config["stitch_gap"] == expected_stitch
assert behavior_config["min_bout_length"] == expected_min_bout

@pytest.mark.parametrize("overwrite", [True, False])
@patch("jabs_postprocess.cli.main.generate_behavior_tables")
Expand Down Expand Up @@ -333,5 +368,8 @@ def test_generate_tables_no_behaviors(
result = runner.invoke(app, cmd_args)

# Assert
assert result.exit_code != 0
assert "Missing option" in result.stdout
assert result.exit_code == 1
assert (
"Must provide either --behavior-config or --behavior options"
in result.stdout
)
Loading