keep improving progress tracking

bioimage-io · Oct 15, 2024 · 420ed8b · 420ed8b
1 parent d0696a5
commit 420ed8b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/bioimageio/download/DownloadTracker.java b/src/main/java/io/bioimage/modelrunner/bioimageio/download/DownloadTracker.java
@@ -33,6 +33,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import io.bioimage.modelrunner.engine.installation.EngineInstall;
+
 /**
  * Class that contains the methods to track the progress downloading files. 
  * The files have to be downloaded to the same folder.
@@ -135,9 +137,13 @@ private DownloadTracker(String folder, TwoParameterConsumer<String, Double> cons
 		this.folder = folder;
 		sizeFiles = new LinkedHashMap<String, Long>();
 		for (String link : links) {
+			String key = folder + File.separator + DownloadModel.getFileNameFromURLString(link);
+			if (consumer.get().get(key + EngineInstall.NBYTES_SUFFIX) != null && consumer.get().get(key + EngineInstall.NBYTES_SUFFIX) != -1) {
+				sizeFiles.put(key, consumer.get().get(key + EngineInstall.NBYTES_SUFFIX).longValue());
+				continue;
+			}
 			try {
-				sizeFiles.put(folder + File.separator + DownloadModel.getFileNameFromURLString(link), 
-						DownloadModel.getFileSize(new URL(link)));
+				sizeFiles.put(key, DownloadModel.getFileSize(new URL(link)));
 			} catch (MalformedURLException e) {
 				throw new IOException("The URL '" + link + "' cannot be found.");
 			}

diff --git a/src/main/java/io/bioimage/modelrunner/engine/installation/EngineInstall.java b/src/main/java/io/bioimage/modelrunner/engine/installation/EngineInstall.java
@@ -577,7 +577,7 @@ public long getBasicDownloadTotalSize() {
 				long engineSize = 0;
 				DeepLearningVersion dlVersion = DeepLearningVersion.fromFile(new File(ee.getValue()));
 				for (String link : dlVersion.getJars()) {
-					String key = DownloadModel.getFileNameFromURLString(link) + NBYTES_SUFFIX;
+					String key = ee.getValue() + File.separator + DownloadModel.getFileNameFromURLString(link) + NBYTES_SUFFIX;
 					long val = DownloadModel.getFileSize(new URL(link));
 					this.consumersMap.get(ee.getValue()).accept(key, (double) val);
 					engineSize += val;
@@ -631,7 +631,7 @@ public static boolean installEngineByCompleteName(String engineDir,
 		if (dlVersion == null)
 			throw new IOException("JDLL does not support any engine compatible with the provided engine directory: "
 					+ engineDir);
-		return installEngine(dlVersion, consumer);
+		return installEngineInDir(dlVersion, engineFileDir.getParentFile().getAbsolutePath(), consumer);
 	}
 
 	/**