From e400c314cf00afdea3bdd0f519089a1b2ec64e76 Mon Sep 17 00:00:00 2001
From: Alan O'Callaghan <alan.ocallaghan@outlook.com>
Date: Tue, 25 Feb 2025 16:42:38 +0000
Subject: [PATCH] Pixel size parsing for 0.5 models

---
 .../ext/instanseg/core/InstanSegModel.java    | 50 ++++++++++++++-----
 1 file changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/main/java/qupath/ext/instanseg/core/InstanSegModel.java b/src/main/java/qupath/ext/instanseg/core/InstanSegModel.java
index 228d372..c42d1b6 100644
--- a/src/main/java/qupath/ext/instanseg/core/InstanSegModel.java
+++ b/src/main/java/qupath/ext/instanseg/core/InstanSegModel.java
@@ -11,6 +11,7 @@
 import qupath.bioimageio.spec.Model;
 import qupath.bioimageio.spec.Resource;
 import qupath.bioimageio.spec.tensor.OutputTensor;
+import qupath.bioimageio.spec.tensor.axes.ScaledAxis;
 import qupath.lib.common.GeneralTools;
 import qupath.lib.images.servers.PixelCalibration;
 
@@ -23,6 +24,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -306,6 +308,10 @@ public static Optional<OutputTensorType> fromString(String value) {
     private static int extractChannelNum(Model model) {
         String axes = getAxesString(model.getInputs().getFirst().getAxes());
         int ind = axes.indexOf("c");
+        // v0.5 models use index axis because channel axes have to be fixed length list of channel names
+        if (ind == -1) {
+            ind = axes.indexOf("i");
+        }
         var shape = model.getInputs().getFirst().getShape();
         if (shape.getShapeStep()[ind] == 1) {
             return ANY_CHANNELS;
@@ -430,19 +436,39 @@ private String getREADMEString(Path path) {
     }
 
     private Optional<Map<String, Double>> getPixelSize() {
-        return getModel().flatMap(model -> {
-            var config = model.getConfig().getOrDefault("qupath", null);
-            if (config instanceof Map configMap) {
-                var axes = (List) configMap.get("axes");
-                String x = String.valueOf(((Map) (axes.get(0))).get("step"));
-                String y = String.valueOf(((Map) (axes.get(1))).get("step"));
-                return Optional.of(Map.of(
-                        "x", Double.valueOf(x),
-                        "y", Double.valueOf(y)
-                ));
+        if (model == null) {
+            return Optional.empty();
+        }
+        if (model.isFormatNewerThan("0.5.0")) {
+            var axes = model.getOutputs().getFirst().getAxes();
+            var xAxis = Arrays.stream(axes)
+                    .filter(a -> a.getID().equals("x"))
+                    .findFirst();
+            var yAxis = Arrays.stream(axes)
+                    .filter(a -> a.getID().equals("y"))
+                    .findFirst();
+            if (xAxis.isEmpty() || yAxis.isEmpty()) {
+                return Optional.empty();
             }
-            return Optional.of(Map.of("x", 1.0, "y", 1.0));
-        });
+            return Optional.of(Map.of(
+                        "x", ((ScaledAxis)(xAxis.get())).getScale(),
+                        "y", ((ScaledAxis)(yAxis.get())).getScale()
+            ));
+        } else {
+            return getModel().flatMap(model -> {
+                var config = model.getConfig().getOrDefault("qupath", null);
+                if (config instanceof Map configMap) {
+                    var axes = (List) configMap.get("axes");
+                    String x = String.valueOf(((Map) (axes.get(0))).get("step"));
+                    String y = String.valueOf(((Map) (axes.get(1))).get("step"));
+                    return Optional.of(Map.of(
+                            "x", Double.valueOf(x),
+                            "y", Double.valueOf(y)
+                    ));
+                }
+                return Optional.of(Map.of("x", 1.0, "y", 1.0));
+            });
+        }
     }
 
     /**