custom model logic

Author: jumabon

src/main/java/fiji/plugin/appose/cellpose/cp3/CellposeAppose.java

@@ -17,6 +17,7 @@
 import java.util.List;
 import java.util.Map;
 
+import javax.management.RuntimeErrorException;
 import javax.swing.JDialog;
 import javax.swing.JProgressBar;
 import javax.swing.WindowConstants;
@@ -69,7 +70,7 @@ public class CellposeAppose extends DynamicCommand implements Initializable
 			"neurips_cellpose_transformer"}, description="Choose CP model to run")
 	private String cp_model = "cyto3"; // cellpose model
 
-	@Parameter(label = "Custom model", description = "Custom model path, overrides the Cellpose model", style="file", required = false)
+	@Parameter(label = "Custom model", description = "Custom model path, overrides the Cellpose model", style="file", required = false, validater = "validateCustomModel")
 	private File custom_model = null;
 
 	@Parameter( label = "Diameter", min="0", description="Average diameter of a cell/nuclei (in pixels)" )
@@ -324,7 +325,9 @@ public < T extends RealType< T > & NativeType< T > > void process( final ImagePl
 		final Map< String, Object > inputs = new HashMap<>();
 		inputs.put( "image", NDArrays.asNDArray( img ) );
 		inputs.put( "use_3D", use3d );
-		inputs.put( "model", cp_model );
+		// return null if custom model
+		inputs.put( "model", ( custom_model == null ) ? cp_model : null );
+		inputs.put( "custom_model", ( custom_model == null ) ? null : custom_model.toString() );
 		inputs.put( "diameter", cell_diameter );
 		inputs.put( "cell_channel", parseChannelChoice( cyto_channel ) );
 		inputs.put( "nuclei_channel", parseChannelChoice( nuclei_channel ));
@@ -583,5 +586,14 @@ public static Integer parseChannelChoice(String str) {
 	    return Integer.parseInt(str);
 	}
 
+	public void validateCustomModel() {
+		if ( custom_model != null) {
+			if ( ! custom_model.exists() ) {
+				IJ.error("The path " + custom_model.toString() + " does not exist !");
+				throw new RuntimeException("The path " + custom_model.toString() + " does not exist !");
+			}
+		}
+	}
+	
 
 }

src/main/resources/cp3.py

@@ -63,6 +63,7 @@ def run_cellpose_v3(img: np.ndarray, kwargs: dict) -> tuple[np.ndarray, np.ndarr
 
     model = models.CellposeModel(
         model_type=kwargs.get('model_name', 'cyto3'),
+        pretrained_model=kwargs.get('custom_model', None),
         gpu=kwargs.get('use_gpu', False),
         device=kwargs.get('device', None)
     )
@@ -133,6 +134,7 @@ def share_as_ndarray(img):
 else:
     file = '../../../sample_data/test.tif'
     input_image = io.imread(file)
+    custom_model = None
     model = 'cyto3'
     diameter = 30
     channels = [0, 1]
@@ -157,6 +159,7 @@ def share_as_ndarray(img):
     input_image,
     kwargs={
         "model_name": model,
+        "custom_model": custom_model,
         "channels": channels,
         "diameter": diameter,
         "use_3D": use_3D,