Refactor models API

benches/bench_stardist.rs

@@ -1,9 +1,9 @@
-use cellcast::models::stardist_2d::{predict_versatile_fluo, warm_up_versatile_fluo};
-use cellcast::models::stardist_3d::{predict_demo, warm_up_demo};
 use criterion::{Criterion, criterion_group, criterion_main};
 use imgal::simulation::blob::logistic_metaballs;
 use ndarray::{Ix2, Ix3, arr2};
 
+use cellcast::models::{StarDist2D, StarDist3D};
+
 const SHAPE_2D: [usize; 2] = [128, 128];
 const SHAPE_3D: [usize; 3] = [64, 128, 128];
 const GPU: bool = true;
@@ -28,11 +28,12 @@ fn bench_stardist_2d(c: &mut Criterion) {
     )
     .unwrap();
     let data = data.into_dimensionality::<Ix2>().unwrap();
-    let mut group = c.benchmark_group("stardist_2d");
-    warm_up_versatile_fluo(GPU);
-    group.bench_function("predict_versatile_fluo", |b| {
+    let mut group = c.benchmark_group("StarDist2D");
+    let sd = StarDist2D::init_fluo(None, GPU);
+    sd.warm_up_fluo();
+    group.bench_function("predict_fluo", |b| {
         b.iter(|| {
-            let _ = predict_versatile_fluo(&data, None, None, None, None, GPU).unwrap();
+            let _ = sd.predict_fluo(&data, None, None, None, None).unwrap();
         });
     });
     group.finish();
@@ -70,12 +71,15 @@ fn bench_stardist_3d(c: &mut Criterion) {
     )
     .unwrap();
     let data = data.into_dimensionality::<Ix3>().unwrap();
-    let mut group = c.benchmark_group("stardist_3d");
+    let mut group = c.benchmark_group("StarDist3D");
     group.sample_size(10);
-    warm_up_demo(GPU);
-    group.bench_function("predict_demo", |b| {
+    let sd = StarDist3D::init_fluo(None, GPU);
+    sd.warm_up_fluo();
+    group.bench_function("predict_fluo", |b| {
         b.iter(|| {
-            let _ = predict_demo(&data, None, None, None, None, None, GPU).unwrap();
+            let _ = sd
+                .predict_fluo(&data, None, None, None, None, None)
+                .unwrap();
         });
     });
     group.finish();

crates/cellcast/src/config/mod.rs

@@ -1 +1,2 @@
 pub(crate) mod backend;
+pub(crate) mod weights;

crates/cellcast/src/config/weights.rs

@@ -0,0 +1,4 @@
+pub const DEMO_3D_URL: &'static str =
+    "https://github.com/uw-loci/cellcast/raw/refs/heads/main/weights/stardist/stardist_3d_demo.bpk";
+pub const VERSATILE_FLUO_2D_URL: &'static str = "https://github.com/uw-loci/cellcast/raw/refs/tags/cellcast-v0.1.2/weights/stardist/stardist_2d_versatile_fluo.bpk";
+pub const VERSATILE_HE_2D_URL: &'static str = "https://github.com/uw-loci/cellcast/raw/refs/tags/cellcast-v0.2.0/weights/stardist/stardist_2d_versatile_he.bpk";

crates/cellcast/src/models/mod.rs

@@ -1,2 +1,5 @@
-pub mod stardist_2d;
-pub mod stardist_3d;
+mod stardist_2d;
+mod stardist_3d;
+
+pub use stardist_2d::StarDist2D;
+pub use stardist_3d::StarDist3D;