increase robustness sending to other process

bioimage-io · Oct 2, 2024 · 65b4ed6 · 65b4ed6
1 parent b04a279
commit 65b4ed6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/JavaWorker.java b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/JavaWorker.java
@@ -87,7 +87,7 @@ private void executeScript(String script, Map<String, Object> inputs) {
 		this.reportLaunch();
 		try {
 			if (script.equals("loadModel")) {
-				pi.loadModel((String) inputs.get("modelFolder"), null);
+				pi.loadModel((String) inputs.get("modelFolder"), (String) inputs.get("modelSource"));
 			} else if (script.equals("inference")) {
 				pi.runFromShmas((List<String>) inputs.get("inputs"), (List<String>) inputs.get("outputs"));
 			} else if (script.equals("close")) {

diff --git a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/PytorchJavaCPPInterface.java b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/PytorchJavaCPPInterface.java
@@ -94,6 +94,8 @@ public class PytorchJavaCPPInterface implements DeepLearningEngineInterface
 	 */
 	private JitModule model;
 
+	private String modelFolder;
+
 	private String modelSource;
 
 	private boolean interprocessing = true;
@@ -157,6 +159,7 @@ private Service getRunner() throws IOException, URISyntaxException {
 	 */
 	@Override
 	public void loadModel(String modelFolder, String modelSource) throws LoadModelException {
+		this.modelFolder = modelFolder;
 		this.modelSource = modelSource;
 		if (interprocessing) {
 			try {
@@ -177,7 +180,8 @@ public void loadModel(String modelFolder, String modelSource) throws LoadModelEx
 
 	private void launchModelLoadOnProcess() throws IOException, InterruptedException {
 		HashMap<String, Object> args = new HashMap<String, Object>();
-		args.put("modelFolder", this.modelSource);
+		args.put("modelFolder", this.modelFolder);
+		args.put("modelSource", this.modelSource);
 		Task task = runner.task("loadModel", args);
 		task.waitFor();
 		if (task.status == TaskStatus.CANCELED)