From c9847a04f9ff2fb47c75ab9221f6c49e7aef0662 Mon Sep 17 00:00:00 2001
From: nanguage <nanguage@yahoo.com>
Date: Tue, 23 Apr 2024 09:41:12 +0200
Subject: [PATCH] allow specify channel

---
 public/plugins/chatbot-extension.imjoy.html | 50 +++++++++++++++++----
 public/plugins/chatbot-utils.imjoy.html     | 10 +++--
 public/plugins/ufish.imjoy.html             | 12 +++--
 src/main.ts                                 |  7 +--
 src/stores/run.ts                           |  8 ++++
 src/views/PredictView.vue                   | 11 +++--
 6 files changed, 75 insertions(+), 23 deletions(-)

diff --git a/public/plugins/chatbot-extension.imjoy.html b/public/plugins/chatbot-extension.imjoy.html
index 9b13594..3941278 100644
--- a/public/plugins/chatbot-extension.imjoy.html
+++ b/public/plugins/chatbot-extension.imjoy.html
@@ -25,6 +25,7 @@
 <script lang="javascript">
 const chatbot_url = "https://bioimage.io/chat/"
 const ufish_url = "https://ufish-team.github.io/"
+//const ufish_url = "http://localhost:5173/"
 const utils_url = ufish_url + "/plugins/chatbot-utils.imjoy.html"
 const py_console_url = "https://nanguage.github.io/web-python-console/"
 
@@ -38,7 +39,10 @@
       if (chatbot) {
         await this.registerExtensions(chatbot.registerExtension)
       } else {
-        chatbot = await api.createWindow({src: chatbot_url, name: "BioImage.IO Chatbot"})
+        chatbot = await api.createWindow({
+          src: chatbot_url, name: "BioImage.IO Chatbot",
+          w: 20, h: 40
+        })
         await this.registerExtensions(chatbot.registerExtension)
       }
     }
@@ -54,6 +58,7 @@
       client = await api.createWindow({
         src: ufish_url,
         name: "ufish",
+        w: 25, h: 8
       });
     }
     return client
@@ -75,6 +80,7 @@
       plugin = await api.createWindow({
         src: py_console_url,
         name: "web-python-console",
+        w: 25, h: 30
       });
     }
     return plugin
@@ -92,13 +98,28 @@
           title: "RunUfish",
           description: "Run U-FISH web client to detect spots in the image",
           properties: {
-          },
-        };
+            "channel": {
+              type: "number",
+              description: "Channel to use",
+              default: 0,
+            },
+            "pThreshold": {
+              type: "number",
+              description: "Threshold for spot detection",
+              default: 0.5,
+            },
+            "viewEnhanced": {
+              type: "boolean",
+              description: "View enhanced image",
+              default: false,
+            },
+          }
+        }
       },
       execute: async (config) => {
         const client = await this.getUfishClient()
         await client.waitReady();
-        await client.runPredict();
+        await client.runPredict(config.channel, config.pThreshold, config.viewEnhanced);
         const output = await client.getOutput();
         this.ufish_spots = output.spots
         return "Done"
@@ -119,18 +140,23 @@
             "diameter": {
               type: "number",
               description: "Diameter of the cells",
-              default: 60,
+              default: 40,
             },
             "model_type": {
               type: "string",
-              description: "Model type, cyto or nuclei",
-              default: "cyto",
+              description: "Model type, nuclei or cyto",
+              default: "nuclei",
             },
             "flow_threshold": {
               type: "number",
               description: "Flow threshold",
               default: 0.4,
             },
+            "channel": {
+              type: "number",
+              description: "Channel to use",
+              default: 2,
+            }
           },
         };
       },
@@ -138,7 +164,15 @@
         const client = await this.getUfishClient()
         const utils = await this.getChatbotUtils()
         const img = await client.getInputImage()
-        await utils.run_cellpose(img, {diameter: config.diameter})
+        await utils.run_cellpose(
+          img,
+          config.channel,
+          {
+            diameter: config.diameter,
+            model_type: config.model_type,
+            flow_threshold: config.flow_threshold,
+          }
+        )
         await utils.compute_mask_coutours()
         const coutours = await utils.get_coutours()
         const kaibu = await api.getWindow("Kaibu")
diff --git a/public/plugins/chatbot-utils.imjoy.html b/public/plugins/chatbot-utils.imjoy.html
index 81ea00b..ba4a00c 100644
--- a/public/plugins/chatbot-utils.imjoy.html
+++ b/public/plugins/chatbot-utils.imjoy.html
@@ -63,15 +63,17 @@
         mask = await run_model(model, img, kwargs)
         self.mask = mask
 
-    async def run_cellpose(self, img, kwargs):
+    async def run_cellpose(self, img, channel, kwargs):
         img = img.astype("float32")
-        if img.ndim == 2:
-            img = img[None, :, :]
-        elif img.ndim == 3:
+        if img.ndim == 3:
             # RGB image
             # swap min channel to first
             min_channel = np.argmin(img.shape)
             img = np.moveaxis(img, min_channel, 0)
+            img = img[channel]
+        elif img.ndim > 3:
+            raise ValueError("Image should be 2D or 3D")
+        img = img[None, :, :]
         mask = await run_model("cellpose", img, kwargs)
         self.mask = mask
 
diff --git a/public/plugins/ufish.imjoy.html b/public/plugins/ufish.imjoy.html
index aba4250..8582037 100644
--- a/public/plugins/ufish.imjoy.html
+++ b/public/plugins/ufish.imjoy.html
@@ -119,6 +119,7 @@
             src="https://kaibu.org/#/app",
             window_id="kaibu-container",
             name="Kaibu",
+            w=38, h=20
         )
         self.image = None  # input image
         self.enhanced_image = None  # enhanced image
@@ -159,10 +160,10 @@
         await self.viewer.view_image(image, name=name)
         return image.shape
 
-    async def process_enhanced(self, enh_img, view_enhanced=True):
+    async def process_enhanced(self, enh_img, view_enhanced=True, pThreshold=0.5):
         if view_enhanced:
             await self.viewer.view_image(enh_img, name="enhanced")
-        df = call_spots_local_maxima(enh_img)
+        df = call_spots_local_maxima(enh_img, intensity_threshold=pThreshold)
         coords = df.values
         self.enhanced_image = enh_img
         self.spots = coords
@@ -178,11 +179,14 @@
     async def get_spots(self):
         return self.spots
 
-    async def scale_image(self):
+    async def scale_image(self, channel=None):
         image = self.image
         if len(image.shape) == 3:
             channel_axis = np.argmin(image.shape)
-            image = image.mean(axis=channel_axis)
+            if channel is not None:
+                image = image.take(channel, axis=channel_axis)
+            else:
+                image = image.mean(axis=channel_axis)
         elif len(image.shape) != 2:
             raise ValueError(
                 f"Image has {len(image.shape)} dimensions. "
diff --git a/src/main.ts b/src/main.ts
index 47d2a88..5159ad6 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -40,8 +40,9 @@ async function createApi() {
     await runStore.waitRunable()
   }
 
-  async function run() {
+  async function run(channel=null, pThreshold=0.5, viewEnhanced=true) {
     const runStore = useRunStore()
+    runStore.setParams(channel, pThreshold, viewEnhanced)
     await runStore.run()
   }
 
@@ -60,10 +61,10 @@ async function createApi() {
     return res
   }
 
-  async function predict(inputImage: object) {
+  async function predict(inputImage: object, channel=null, pThreshold=0.5, viewEnhanced=true) {
     await waitRunable()
     await setInputImage(inputImage, "input")
-    await run()
+    await run(channel, pThreshold, viewEnhanced)
     return await getOutput()
   }
 
diff --git a/src/stores/run.ts b/src/stores/run.ts
index f3a0359..d999656 100644
--- a/src/stores/run.ts
+++ b/src/stores/run.ts
@@ -46,6 +46,9 @@ export const useRunStore = defineStore("run", {
     imageUrl: null as string | null,
     fetchedImage: null as object | null,
     fetchGetable: false,
+    channel: null as (number | null),
+    pThreshold: 0.5,
+    viewEnhanced: true,
   }),
   actions: {
     async waitRunable() {
@@ -87,6 +90,11 @@ export const useRunStore = defineStore("run", {
     async getFetchedImage() {
       await waitFetchGetable();
       return this.fetchedImage;
+    },
+    setParams(channel: number | null, pThreshold: number, viewEnhanced: boolean) {
+      this.channel = channel;
+      this.pThreshold = pThreshold;
+      this.viewEnhanced = viewEnhanced;
     }
   }
 }) 
\ No newline at end of file
diff --git a/src/views/PredictView.vue b/src/views/PredictView.vue
index b9563c7..f495532 100644
--- a/src/views/PredictView.vue
+++ b/src/views/PredictView.vue
@@ -106,7 +106,7 @@ export default {
     })
 
     function checkInputShape(shape: number[]) {
-      if (shape.length === 3 && shape[2] === 3) {
+      if (shape.length === 3 && ((shape[2] === 3) || (shape[2] === 4))) {
         runInfoText.value = `Image loaded, shape: ${shape}. Will treat it as an RGB image.`
         hasError.value = false
       } else if (shape.length !== 2) {
@@ -120,8 +120,11 @@ export default {
 
     async function run() {
       running.value = true
+      const channel = runStore.channel
+      const pThreshold = runStore.pThreshold
+      const viewEnhanced = runStore.viewEnhanced
       try {
-        const sImg = await plugin.value.scale_image()
+        const sImg = await plugin.value.scale_image(channel)
         const f32data = new Float32Array(sImg._rvalue)
         const input = new ort.Tensor(
         sImg._rdtype, f32data, sImg._rshape);
@@ -132,8 +135,8 @@ export default {
           _rshape: modelOut.dims,
           _rvalue: (modelOut.data as Float32Array).buffer
         }
-        const viewEnhanced = !isPluginMode()
-        const [enhBytes, coords, numSpots] = await plugin.value.process_enhanced(outImg, viewEnhanced)
+        const [enhBytes, coords, numSpots] = await plugin.value.process_enhanced(
+          outImg, viewEnhanced, pThreshold)
         output.value = {
           enhanced: enhBytes,
           coords: coords