Fix model downloading bug

Previously it would show a downloading notification twice, then try too quickly to run the model before it was downloaded - and fail.

This also fixes the output channel checkboxes.

src/main/java/qupath/ext/instanseg/core/DetectionMeasurer.java

@@ -14,6 +14,12 @@
 import java.util.Collection;
 import java.util.List;
 
+/**
+ * Helper class for adding measurements to InstanSeg detections.
+ * <p>
+ * Note that this is inherently limited to 'small' detections, where the entire ROI can be loaded into memory.
+ * It does not support measuring arbitrarily large regions at a high resolution.
+ */
 public class DetectionMeasurer {
 
     private static final Logger logger = LoggerFactory.getLogger(DetectionMeasurer.class);

src/main/java/qupath/ext/instanseg/core/InstanSeg.java

@@ -54,7 +54,8 @@ public class InstanSeg {
     private final TaskRunner taskRunner;
     private final Class<? extends PathObject> preferredOutputClass;
 
-    // This was previously an adjustable parameter, but it's now fixed at 1 because we handle overlaps differently
+    // This was previously an adjustable parameter, but it's now fixed at 1 because we handle overlaps differently.
+    // However we might want to reinstate it, possibly as a proportion of the padding amount.
     private final int boundaryThreshold = 1;
 
     private InstanSeg(Builder builder) {
@@ -119,6 +120,7 @@ public InstanSegResults detectObjects(ImageData<BufferedImage> imageData, Collec
     private void validateImageAndObjectsOrThrow(ImageData<BufferedImage> imageData, Collection<? extends PathObject> pathObjects) {
         Objects.requireNonNull(imageData, "No imageData available");
         Objects.requireNonNull(pathObjects, "No objects available");
+        // TODO: Consider if there are use cases where it is worthwhile to provide objects that are not in the hierarchy
         var hierarchy = imageData.getHierarchy();
         if (pathObjects.stream().anyMatch(p -> !PathObjectTools.hierarchyContainsObject(hierarchy, p))) {
             throw new IllegalArgumentException("Objects must be contained in the image hierarchy!");

src/main/java/qupath/ext/instanseg/core/InstanSegModel.java

@@ -9,7 +9,6 @@
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import qupath.lib.common.GeneralTools;
-import qupath.lib.gui.UserDirectoryManager;
 import qupath.lib.images.servers.PixelCalibration;
 
 import java.io.IOException;

src/main/java/qupath/ext/instanseg/ui/InstanSegController.java

@@ -76,6 +76,7 @@
 import java.util.Optional;
 import java.util.ResourceBundle;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ForkJoinPool;
@@ -436,55 +437,91 @@ private void configureModelChoices() {
             handleModelDirectory(n);
 //            addRemoteModels(modelChoiceBox.getItems());
         });
-        modelChoiceBox.getSelectionModel().selectedItemProperty().addListener((v, o, n) -> {
-            if (n == null) {
-                return;
-            }
-            var modelDir = getModelDirectory().orElse(null);
-            boolean isDownloaded = modelDir != null && n.isDownloaded(modelDir);
-            if (!isDownloaded || qupath.getImageData() == null) {
-                return;
-            }
-            var numChannels = n.getNumChannels();
-            if (qupath.getImageData().isBrightfield() && numChannels.isPresent() && numChannels.get() != InstanSegModel.ANY_CHANNELS) {
-                comboChannels.getCheckModel().clearChecks();
-                comboChannels.getCheckModel().checkIndices(0, 1, 2);
-            }
-            // Handle output channels
-            var nOutputs = n.getOutputChannels().orElse(1);
-            checkComboOutputs.getItems().setAll(InstanSegOutput.getOutputsForChannelCount(nOutputs));
-            checkComboOutputs.getCheckModel().checkAll();
-        });
-        downloadButton.setOnAction(e -> downloadModel());
+        modelChoiceBox.getSelectionModel().selectedItemProperty().addListener((v, o, n) -> refreshModelChoice());
+        downloadButton.setOnAction(e -> downloadSelectedModelAsync());
         WebView webView = WebViews.create(true);
         PopOver infoPopover = new PopOver(webView);
         infoButton.setOnAction(e -> {
             parseMarkdown(modelChoiceBox.getSelectionModel().getSelectedItem(), webView, infoButton, infoPopover);
         });
     }
 
-    private void downloadModel() {
-        try (var pool = ForkJoinPool.commonPool()) {
-            pool.execute(() -> {
-                try {
-                    var modelDir = getModelDirectory().orElse(null);
-                    if (modelDir == null || !Files.exists(modelDir)) {
-                        Dialogs.showErrorMessage(resources.getString("title"),
-                                resources.getString("ui.model-directory.choose-prompt"));
-                        return;
-                    }
-                    var model = modelChoiceBox.getSelectionModel().getSelectedItem();
-                    Dialogs.showInfoNotification(resources.getString("title"),
-                            String.format(resources.getString("ui.popup.fetching"), model.getName()));
-                    model.download(modelDir);
-                    Dialogs.showInfoNotification(resources.getString("title"),
-                            String.format(resources.getString("ui.popup.available"), model.getName()));
-                    needsUpdating.set(!needsUpdating.get());
-                } catch (IOException ex) {
-                    Dialogs.showErrorNotification(resources.getString("title"), resources.getString("error.downloading"));
-                }
+    /**
+     * Make UI changes based on the selected model.
+     * This may be called when the selected model is changed, or an existing model is downloaded.
+     */
+    private void refreshModelChoice() {
+        var model = modelChoiceBox.getSelectionModel().getSelectedItem();
+        if (model == null)
+            return;
+
+        var modelDir = getModelDirectory().orElse(null);
+        boolean isDownloaded = modelDir != null && model.isDownloaded(modelDir);
+        if (!isDownloaded || qupath.getImageData() == null) {
+            return;
+        }
+        var numChannels = model.getNumChannels();
+        if (qupath.getImageData().isBrightfield() && numChannels.isPresent() && numChannels.get() != InstanSegModel.ANY_CHANNELS) {
+            comboChannels.getCheckModel().clearChecks();
+            comboChannels.getCheckModel().checkIndices(0, 1, 2);
+        }
+        // Handle output channels
+        var nOutputs = model.getOutputChannels().orElse(1);
+        checkComboOutputs.getCheckModel().clearChecks();
+        checkComboOutputs.getItems().setAll(InstanSegOutput.getOutputsForChannelCount(nOutputs));
+        checkComboOutputs.getCheckModel().checkAll();
+    }
+
+    /**
+     * Try to download the currently-selected model in another thread.
+     * @return
+     */
+    private CompletableFuture<InstanSegModel> downloadSelectedModelAsync() {
+        var model = modelChoiceBox.getSelectionModel().getSelectedItem();
+        if (model == null) {
+            return CompletableFuture.completedFuture(null);
+        }
+        return downloadModelAsync(model);
+    }
+
+    /**
+     * Try to download the specified model in another thread.
+     * @param model
+     * @return
+     */
+    private CompletableFuture<InstanSegModel> downloadModelAsync(InstanSegModel model) {
+        var modelDir = getModelDirectory().orElse(null);
+        if (modelDir == null || !Files.exists(modelDir)) {
+            Dialogs.showErrorMessage(resources.getString("title"),
+                    resources.getString("ui.model-directory.choose-prompt"));
+            return CompletableFuture.completedFuture(null);
+        }
+        return CompletableFuture.supplyAsync(() -> downloadModel(model, modelDir), ForkJoinPool.commonPool());
+    }
+
+    /**
+     * Try to download the specified model to the given directory in the current thread.
+     * @param model
+     * @param modelDir
+     * @return
+     */
+    private InstanSegModel downloadModel(InstanSegModel model, Path modelDir) {
+        Objects.requireNonNull(modelDir);
+        Objects.requireNonNull(model);
+        try {
+            Dialogs.showInfoNotification(resources.getString("title"),
+                    String.format(resources.getString("ui.popup.fetching"), model.getName()));
+            model.download(modelDir);
+            Dialogs.showInfoNotification(resources.getString("title"),
+                    String.format(resources.getString("ui.popup.available"), model.getName()));
+            FXUtils.runOnApplicationThread(() -> {
+                needsUpdating.set(!needsUpdating.get());
+                refreshModelChoice();
             });
+        } catch (IOException ex) {
+            Dialogs.showErrorNotification(resources.getString("title"), resources.getString("error.downloading"));
         }
+        return model;
     }
 
     private static void parseMarkdown(InstanSegModel model, WebView webView, Button infoButton, PopOver infoPopover) {
@@ -676,41 +713,57 @@ void restart() {
 
     @FXML
     private void runInstanSeg() {
+        runInstanSeg(modelChoiceBox.getSelectionModel().getSelectedItem());
+    }
+
+    private void runInstanSeg(InstanSegModel model) {
+        if (model == null) {
+            Dialogs.showErrorNotification(resources.getString("title"), resources.getString("ui.error.no-model"));
+            return;
+        }
+        var imageData = qupath.getImageData();
+        if (imageData == null) {
+            Dialogs.showErrorNotification(resources.getString("title"), resources.getString("error.no-imagedata"));
+            return;
+        }
+
         if (!PytorchManager.hasPyTorchEngine()) {
             if (!Dialogs.showConfirmDialog(resources.getString("title"), resources.getString("ui.pytorch"))) {
                 Dialogs.showWarningNotification(resources.getString("title"), resources.getString("ui.pytorch-popup"));
                 return;
             }
         }
-        ImageServer<?> server = qupath.getImageData().getServer();
-        List<ChannelSelectItem> selectedChannels = comboChannels
-                .getCheckModel().getCheckedItems()
-                .stream()
-                .filter(Objects::nonNull)
-                .toList();
 
-
-        var model = modelChoiceBox.getSelectionModel().getSelectedItem();
         var modelPath = getModelDirectory().orElse(null);
         if (modelPath == null) {
             Dialogs.showErrorNotification(resources.getString("title"), resources.getString("ui.model-directory.choose-prompt"));
             return;
         }
+
         if (!model.isDownloaded(modelPath)) {
             if (!Dialogs.showYesNoDialog(resources.getString("title"), resources.getString("ui.model-popup")))
                 return;
-            Dialogs.showInfoNotification(resources.getString("title"), String.format(resources.getString("ui.popup.fetching"), model.getName()));
-            downloadModel();
-            if (!model.isDownloaded(modelPath)) {
-                Dialogs.showErrorNotification(resources.getString("title"), String.format(resources.getString("error.localModel")));
-                return;
-            }
+            downloadModelAsync(model)
+                    .thenAccept((InstanSegModel suppliedModel) -> {
+                        if (suppliedModel == null || !suppliedModel.isDownloaded(modelPath)) {
+                            Dialogs.showErrorNotification(resources.getString("title"), String.format(resources.getString("error.localModel")));
+                        } else {
+                            runInstanSeg(suppliedModel);
+                        }
+                    });
+            return;
         }
 
+        ImageServer<?> server = imageData.getServer();
+        List<ChannelSelectItem> selectedChannels = comboChannels
+                .getCheckModel().getCheckedItems()
+                .stream()
+                .filter(Objects::nonNull)
+                .toList();
         int imageChannels = selectedChannels.size();
         var modelChannels = model.getNumChannels();
         if (modelChannels.isEmpty()) {
-            Dialogs.showErrorNotification(resources.getString("title"), resources.getString("error.fetching"));
+            Dialogs.showErrorNotification(resources.getString("title"), resources.getString("ui.error.model-not-downloaded"));
             return;
         }
 

src/main/java/qupath/ext/instanseg/ui/InstanSegExtension.java

@@ -26,9 +26,9 @@ public class InstanSegExtension implements QuPathExtension, GitHubProject {
 
 	private static final String EXTENSION_NAME = "InstanSeg";
 
-	private static final String EXTENSION_DESCRIPTION = "Use the InstanSeg deep learning model in QuPath";
+	private static final String EXTENSION_DESCRIPTION = "Use InstanSeg deep learning models for inference in QuPath";
 
-	private static final Version EXTENSION_QUPATH_VERSION = Version.parse("v0.5.0");
+	private static final Version EXTENSION_QUPATH_VERSION = Version.parse("v0.6.0");
 
 	private static final GitHubRepo EXTENSION_REPOSITORY = GitHubRepo.create(
 			EXTENSION_NAME, "qupath", "qupath-extension-instanseg");

src/main/resources/qupath/ext/instanseg/ui/strings.properties

@@ -97,7 +97,7 @@ ui.stop-tasks = Stop all running tasks?
 
 ## Errors
 error.window = Error initializing InstanSeg window.\nAn internet connection is required when running for the first time.
-error.no-imagedata = Cannot run InstanSeg without ImageData.
+error.no-imagedata = Cannot run InstanSeg without an image available.
 error.downloading = Error downloading files
 error.querying-local = Error querying local files
 error.localModel = Can't find file in user model directory