Minor fixes to inference workflow (#360)

In this PR, the following are updated:

This PR sets integral refinement to None by default when running on
Apple MPS (unsupported op), replacing the previous CPU fallback that
caused slowdowns. The pipeline now stays on MPS unless the user
explicitly re-enables refinement. It also fixes a test issue where
output path was passed as a directory name instead of a file, triggering
a “frame error skipped” condition; tests now pass a valid video path.

In this PR, we also implement the suggested fix
[here](#357) to gracefully
handle unicode encoding errors in the logger. (Related issue: #357)

Author: gitttt-1234

sleap_nn/__init__.py

@@ -1,6 +1,7 @@
 """Main module for sleap_nn package."""
 
 import os
+import sys
 from loguru import logger
 
 # Get RANK for distributed training
@@ -24,9 +25,24 @@ def _should_log(record):
 # Remove default handler and add custom one
 logger.remove()
 
+
+def _safe_print(msg):
+    """Print with fallback for encoding errors."""
+    try:
+        print(msg, end="")
+    except UnicodeEncodeError:
+        # Fallback: replace unencodable characters with '?'
+        print(
+            msg.encode(sys.stdout.encoding, errors="replace").decode(
+                sys.stdout.encoding
+            ),
+            end="",
+        )
+
+
 # Add logger with the custom filter
 logger.add(
-    lambda msg: print(msg, end=""),
+    _safe_print,
     level="DEBUG",
     filter=_should_log,
     format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {name}:{function}:{line} | {message}",

sleap_nn/predict.py

@@ -358,9 +358,9 @@ def run_inference(
         if integral_refinement is not None and device == "mps":  # TODO
             # kornia/geometry/transform/imgwarp.py:382: in get_perspective_transform. NotImplementedError: The operator 'aten::_linalg_solve_ex.result' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
             logger.info(
-                "Integral refinement is not supported with MPS device. Using CPU."
+                "Integral refinement is not supported with MPS accelerator. Setting integral refinement to None."
             )
-            device = "cpu"  # not supported with mps
+            integral_refinement = None
 
         logger.info(f"Using device: {device}")
 

tests/test_cli.py

@@ -7,6 +7,7 @@
 import subprocess
 from sleap_nn.predict import run_inference
 import sleap_io as sio
+import torch
 
 
 @pytest.fixture
@@ -164,6 +165,8 @@ def test_track_command(
         f"{tmp_path}/test.slp",
         "--frames",
         "0-99",
+        "--device",
+        "cpu",
     ]
     result = subprocess.run(cmd, check=True, capture_output=True, text=True)
     assert Path(f"{tmp_path}/test.slp").exists()
@@ -182,6 +185,7 @@ def test_eval_command(
         make_labels=True,
         max_instances=6,
         output_path=f"{tmp_path}/test.slp",
+        device="cpu" if torch.backends.mps.is_available() else "auto",
     )
     cmd = [
         "uv",

tests/test_evaluation.py

@@ -504,6 +504,7 @@ def test_evaluator_main(
         make_labels=True,
         max_instances=6,
         output_path=f"{tmp_path}/test.slp",
+        device="cpu" if torch.backends.mps.is_available() else "auto",
     )
 
     import subprocess