Modify inference and centroid model to allow Bonsai ingestion (#850)

* Add WIP frozen graph saving to inference model

* Add option for optimal grouping in multiclass peaks

* Add centroid only model

* Add model save tests

* Update save tests

* Lint peak finding

* Call unwrapped find_global_peaks in tests to raise codecov

* Call unwrapped find_local_peaks in tests to raise codecov

* Call unwrapped CentroidCrop.call in tests to raise codecov

* Decorate offset peak finding, add topdown save tests

* Add return_crops flag to centroid crop layer

* Update bonsai model export, add high level export, cli func

* Add predictor export tests

* Expose high level exporter

* Add entry point for cli export

* Add tracing for resizing pad_to_stride

* Call unwrapped pad_to_stride in tests to raise codecov

* Lint resizing test

Co-authored-by: roomrys <[email protected]>

setup.py

@@ -58,6 +58,7 @@ def get_requirements(require_name=None):
             "sleap-track=sleap.nn.inference:main",
             "sleap-inspect=sleap.info.labels:main",
             "sleap-diagnostic=sleap.diagnostic:main",
+            "sleap-export=sleap.nn.inference:export_cli",
         ],
     },
     python_requires=">=3.6",

sleap/__init__.py

@@ -14,7 +14,7 @@
 import sleap.nn
 from sleap.nn.data import pipelines
 from sleap.nn import inference
-from sleap.nn.inference import load_model
+from sleap.nn.inference import load_model, export_model
 from sleap.nn.system import use_cpu_only, disable_preallocation
 from sleap.nn.system import summary as system_summary
 from sleap.nn.config import TrainingJobConfig, load_config

sleap/nn/data/resizing.py

@@ -30,6 +30,7 @@ def find_padding_for_stride(
     return pad_bottom, pad_right
 
 
+@tf.function
 def pad_to_stride(image: tf.Tensor, max_stride: int) -> tf.Tensor:
     """Pad an image to meet a max stride constraint.