Merge pull request #8 from imglib/newcpsammodels

Newcpsammodels

Author: tinevez

src/main/java/net/imglib2/cellpose/Cellpose4BuiltinModels.java

@@ -0,0 +1,78 @@
+/*-
+ * #%L
+ * Running Cellpose 3 and 4 from Java with Appose, using ImgLib2 data structure.
+ * %%
+ * Copyright (C) 2026 Appose developpers
+ * %%
+ * Redistribution and use in source and binary forms, with or without modification,
+ * are permitted provided that the following conditions are met:
+ * 
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ *    list of conditions and the following disclaimer.
+ * 
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ *    this list of conditions and the following disclaimer in the documentation
+ *    and/or other materials provided with the distribution.
+ * 
+ * 3. Neither the name of the ImgLib2 nor the names of its contributors
+ *    may be used to endorse or promote products derived from this software without
+ *    specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
+ * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+ * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+ * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
+ * OF THE POSSIBILITY OF SUCH DAMAGE.
+ * #L%
+ */
+package net.imglib2.cellpose;
+
+public enum Cellpose4BuiltinModels
+{
+	// General  models
+	CPSAM( "cpsam", "Cellpose-SAM model. General model, works well on most data." ),
+	CPSAMV2( "cpsam_v2", "Cellpose-SAM model released in June 2026. It handles better low contrast regions" );
+	//CPDINO( "cpdino", "CellposeDINO model, based on DINO backbone. For RGB images."),
+	//CPDINOVITB( "cpdino-vitb", "CellposeDINO smaller model, based on DINO backbone. For RGB images.");
+	
+	private final String name;
+
+	private final String description;
+
+	Cellpose4BuiltinModels( final String name, final String description )
+	{
+		this.name = name;
+		this.description = description;
+	}
+
+	public String modelName()
+	{
+		return name;
+	}
+
+	public String description()
+	{
+		return description;
+	}
+
+	@Override
+	public String toString()
+	{
+		return name;
+	}
+
+	/**
+	 * Gets a tooltip string combining the model name and description
+	 * 
+	 * @return Formatted string for tooltip display
+	 */
+	public String getTooltip()
+	{
+		return String.format( "%s: %s", name, description );
+	}
+}

src/main/java/net/imglib2/cellpose/Cellpose4Parameters.java

@@ -35,21 +35,24 @@
 import java.util.Map;
 
 import net.imglib2.appose.ShmImg;
+import net.imglib2.cellpose.Cellpose3Parameters.Builder;
 import net.imglib2.type.NativeType;
 import net.imglib2.type.numeric.IntegerType;
 import net.imglib2.type.numeric.RealType;
 import net.imglib2.type.numeric.integer.UnsignedByteType;
 
 public class Cellpose4Parameters extends CellposeParameters
 {
-
+	public final Cellpose4BuiltinModels buitInModel;
+	
 	public final Integer chan0; // as Integer so that it can be null
 
 	public final Integer chan1;
 
 	public final Integer chan2;
 
 	private Cellpose4Parameters(
+			final Cellpose4BuiltinModels buitInModel,
 			final Integer chan0,
 			final Integer chan1,
 			final Integer chan2,
@@ -75,6 +78,7 @@ private Cellpose4Parameters(
 				cellProbThreshold, useGpu, minSize, anisotropy,
 				stitchThreshold, resample, tileOverlap, computeFlows,
 				flow3dSmooth, nIter, torchVersion );
+		this.buitInModel = buitInModel;
 		this.chan0 = chan0;
 		this.chan1 = chan1;
 		this.chan2 = chan2;
@@ -87,7 +91,10 @@ public < T extends RealType< T > & NativeType< T >, R extends IntegerType< R > &
 			final ShmImg< R > outputLabels,
 			final ShmImg< UnsignedByteType > outputFlows )
 	{
+		
 		final Map< String, Object > inputs = super.toApposeMap( input, axisInfo, outputLabels, outputFlows );
+		final boolean isBuiltInModel = customModel == null || customModel.equals( "" );
+		inputs.put( "model_name", isBuiltInModel ? buitInModel.modelName() : null );
 
 		final long nChannels = axisInfo.nChannels( input );
 		inputs.put( "n_channels", nChannels );
@@ -106,6 +113,13 @@ public static Builder builder()
 	// Builder class for fluent construction
 	public static class Builder extends CellposeParameters.Builder< Builder >
 	{
+		private Cellpose4BuiltinModels model = Cellpose4BuiltinModels.CPSAMV2;
+		
+		public Builder model( final Cellpose4BuiltinModels model )
+		{
+			this.model = model;
+			return this;
+		}
 
 		private Integer chan0 = 0;  // to have one selected by default
 
@@ -135,7 +149,7 @@ public Builder chan2( final Integer chan2 )
 		public Cellpose4Parameters build()
 		{
 			return new Cellpose4Parameters(
-					chan0, chan1, chan2, customModel, diameter, do3D, normalize,
+					model, chan0, chan1, chan2, customModel, diameter, do3D, normalize,
 					flowThreshold, cellProbThreshold, useGpu, minSize,
 					anisotropy, stitchThreshold, resample, tileOverlap,
 					computeFlows, flow3dSmooth, nIter, torchVersion );

src/main/resources/cp4_init.py

@@ -50,17 +50,18 @@
 # load arguments and input from Appose task
 if appose_mode:
     use_gpu: bool = globals()['use_gpu']
-    selected_model = "cpsam" if custom_model is None else custom_model
+    selected_model = model_name if custom_model is None else custom_model
 else:
     custom_model = None
+    model_name = "cpsam_v2"
     use_gpu = False
 
 use_gpu, device = get_torch_device(use_gpu)
 
 task.update(
     current = 1,
     maximum= 2,
-    message=f"CP4: Start Cellpose (device={device}): deploy model {selected_model if selected_model else custom_model}"
+    message=f"CP4: Start Cellpose (device={device}): deploy model {selected_model}"
 )
 
 model = models.CellposeModel(