add more methods to model API

segment-anything-models-java · Aug 7, 2024 · d299a3a · d299a3a
1 parent de7f376
commit d299a3a
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 0 deletions.
diff --git a/src/main/java/ai/nets/samj/install/EfficientSamEnvManager.java b/src/main/java/ai/nets/samj/install/EfficientSamEnvManager.java
@@ -38,6 +38,7 @@
 
 import org.apache.commons.compress.archivers.ArchiveException;
 
+import ai.nets.samj.gui.tools.Files;
 import io.bioimage.modelrunner.apposed.appose.Mamba;
 import io.bioimage.modelrunner.apposed.appose.MambaInstallException;
 
@@ -424,4 +425,22 @@ public String getModelWeigthsName() {
 	public String getModelWeigthPath() {
 		return Paths.get(this.path, "envs", ESAM_ENV_NAME, ESAM_NAME, "weights", ESAM_SMALL_WEIGHTS_NAME).toAbsolutePath().toString();
 	}
+
+	@Override
+	public boolean checkEverythingInstalled() {
+		if (!this.checkMambaInstalled()) return false;
+
+		if (!this.checkSAMDepsInstalled()) return false;
+
+		if (!this.checkEfficientSAMPackageInstalled()) return false;
+
+		if (!this.checkModelWeightsInstalled()) return false;
+
+		return true;
+	}
+
+	@Override
+	public void uninstall() {
+		Files.deleteFolder(new File(this.getModelEnv()));
+	}
 }
diff --git a/src/main/java/ai/nets/samj/install/EfficientViTSamEnvManager.java b/src/main/java/ai/nets/samj/install/EfficientViTSamEnvManager.java
@@ -41,6 +41,7 @@
 
 import org.apache.commons.compress.archivers.ArchiveException;
 
+import ai.nets.samj.gui.tools.Files;
 import ai.nets.samj.models.EfficientViTSamJ;
 import io.bioimage.modelrunner.apposed.appose.Mamba;
 import io.bioimage.modelrunner.apposed.appose.MambaInstallException;
@@ -498,4 +499,25 @@ public String getModelWeigthsName() {
 	public String getModelWeigthPath() {
 		return Paths.get(this.path, "envs", EVITSAM_ENV_NAME, EVITSAM_NAME, "weights", modelType + ".pt").toAbsolutePath().toString();
 	}
+
+	@Override
+	public boolean checkEverythingInstalled() {
+		if (!this.checkMambaInstalled()) return false;
+
+		if (!this.checkSAMDepsInstalled()) return false;
+
+		if (!this.checkEfficientViTSAMPackageInstalled()) return false;
+
+		if (!this.checkModelWeightsInstalled()) return false;
+
+		return true;
+	}
+
+	@Override
+	public void uninstall() {
+		if (new File(this.getModelWeightsPath()).getParentFile().list().length != 1)
+			Files.deleteFolder(new File(this.getModelWeightsPath()));
+		else
+			Files.deleteFolder(new File(this.getModelEnv()));
+	}
 }
diff --git a/src/main/java/ai/nets/samj/install/Sam2EnvManager.java b/src/main/java/ai/nets/samj/install/Sam2EnvManager.java
@@ -35,6 +35,7 @@
 
 import org.apache.commons.compress.archivers.ArchiveException;
 
+import ai.nets.samj.gui.tools.Files;
 import ai.nets.samj.models.Sam2;
 import io.bioimage.modelrunner.apposed.appose.Mamba;
 import io.bioimage.modelrunner.apposed.appose.MambaInstallException;
@@ -384,4 +385,23 @@ public String getModelWeigthPath() {
 			return Paths.get(path, "envs", SAM2_ENV_NAME, SAM2_NAME, "weights", String.format("sam2_hiera_%s.pt", modelType)).toAbsolutePath().toString();
 		}
 	}
+
+	@Override
+	public boolean checkEverythingInstalled() {
+		if (!this.checkMambaInstalled()) return false;
+
+		if (!this.checkSAMDepsInstalled()) return false;
+
+		if (!this.checkModelWeightsInstalled()) return false;
+
+		return true;
+	}
+
+	@Override
+	public void uninstall() {
+		if (new File(this.getModelWeightsPath()).getParentFile().list().length != 1)
+			Files.deleteFolder(new File(this.getModelWeightsPath()));
+		else
+			Files.deleteFolder(new File(this.getModelEnv()));
+	}
 }