support pass parameters to cellpose

bioimage-io · Dec 19, 2023 · e656e15 · e656e15
1 parent 41e2365
commit e656e15
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 22 deletions.
diff --git a/src/components/AdvanceSetting.vue b/src/components/AdvanceSetting.vue
@@ -1,5 +1,5 @@
 <template>
-  <div>
+  <div class="advance-setting-container">
     <div class="server-setting">
       <h3>Settings for server</h3>
       <label class="field-label">Server URL</label>
@@ -61,6 +61,10 @@
 </template>
 
 <style scoped>
+.advance-setting-container {
+  padding-left: 20px;
+  padding-bottom: 20px;
+}
 .field-label {
   display: block;
   margin-bottom: 5px;

diff --git a/src/components/ModelParameters.vue b/src/components/ModelParameters.vue
@@ -0,0 +1,48 @@
+<template>
+  <div v-if="parameters && (parameters.length > 0)">
+    <HideContainer :summary="'Model parameters'">
+      <OverlayContainer :open="overlay">
+        <div class="model-params-container">
+          <div v-for="(paramGroup, index) in parameters" :key="index">
+            <h3>{{ paramGroup.name }}</h3>
+            <div v-for="(param, index2) in paramGroup.parameters" :key="index">
+              <label class="field-label">{{ param.name }}</label>
+              <InputNumber showButtons v-if="param.type === 'number'" v-model="parameterValues[param.name]" />
+              <Dropdown :options="param.enum" v-else-if="param.enum" v-model="parameterValues[param.name]" />
+              <InputText v-else="param.type === 'string'" v-model="parameterValues[param.name]" />
+            </div>
+          </div>
+        </div>
+      </OverlayContainer>
+    </HideContainer>
+  </div>
+</template>
+
+<style scoped>
+.field-label {
+  display: block;
+  margin-bottom: 5px;
+  font-weight: bold;
+}
+.model-params-container {
+  padding-left: 20px;
+  padding-bottom: 20px;
+}
+</style>
+
+<script>
+import HideContainer from './HideContainer.vue';
+import OverlayContainer from './OverlayContainer.vue';
+
+export default {
+  props: {
+    parameters: Array,
+    parameterValues: Object,
+    overlay: Boolean,
+  },
+  components: {
+    HideContainer,
+    OverlayContainer
+},
+}
+</script>
diff --git a/src/components/ModelSelect.vue b/src/components/ModelSelect.vue
@@ -1,6 +1,6 @@
 <template>
   <div>
-    <Dropdown v-model="currentModel" :options="modelList" option-label="name">
+    <Dropdown :disabled="!open" v-model="currentModel" :options="modelList" option-label="name">
       <template #value="slotProps">
         <div v-if="slotProps.value">
           <div>
@@ -25,6 +25,10 @@
 <script>
 export default {
   props: {
+    open: {
+      type: Boolean,
+      default: false
+    },
     additionalModels: {
       type: Array,
       default: () => []

diff --git a/src/components/RunPanel.vue b/src/components/RunPanel.vue
@@ -3,7 +3,7 @@
     <div class="model-selection">
         <div class="model-selection-line">
             <div class="model-selection-label">Model</div>
-            <ModelSelect @model-selected="handleModelChange" :additional-models="additionalModels" />
+            <ModelSelect :open="this.modelSelectEnable" @model-selected="handleModelChange" :additional-models="additionalModels" />
         </div>
         <div class="model-selection-tips">ℹ️  Please visit <a href="https://bioimage.io/#/" target="_blank">bioimage.io</a> to view detailed information about the model.</div>
     </div>
@@ -28,6 +28,11 @@
         Show reference output
       </Button>
     </div>
+    <ModelParameters
+      :parameters="this.additionalParameters"
+      :parameterValues="this.additionalParameterValues"
+      :overlay="waiting"
+    />
     <HideContainer :summary="'Advanced settings'">
       <OverlayContainer :open="waiting">
         <AdvanceSetting
@@ -106,6 +111,7 @@ import OverlayContainer from "./OverlayContainer.vue";
 import AdvanceSetting from "./AdvanceSetting.vue";
 import LoadingAnimation from "./LoadingAnimation.vue";
 import ModelSelect from "./ModelSelect.vue";
+import ModelParameters from "./ModelParameters.vue";
 import { ModelRunner } from "../modelRun";
 import { ImagejJsController } from "../viewerControl";
 
@@ -129,11 +135,13 @@ export default {
         buttonEnabledRun: false,
         buttonEnabledInput: false,
         buttonEnabledOutput: false,
+        modelSelectEnable: false,
         tileSizes: { x: 0, y: 0, z: 0 },
         tileOverlaps: { x: 0, y: 0, z: 0 },
         runner: null,
-        serverUrl: "https://hypha.bioimage.io",
+        serverUrl: "https://ai.imjoy.io",
         additionalModels: [],
+        additionalParameterValues: {}
     }),
     computed: {
         infoColor() {
@@ -151,29 +159,37 @@ export default {
             return this.runner.modelTritonConfig;
         },
         inputMinShape() {
-            if (!this.runner) {
+            if (!this?.runner?.rdf) {
                 return {};
             }
             return this.runner.getInputMinShape();
         },
         inputMaxShape() {
-            if (!this.runner) {
+            if (!this?.runner?.rdf) {
                 return {};
             }
             return this.runner.getInputMaxShape();
         },
         inputSpec() {
-            if (!this.runner) {
+            if (!this?.runner?.rdf) {
                 return {};
             }
             return this.runner.rdf.inputs[0];
         },
         outputSpec() {
-            if (!this.runner) {
+            if (!this?.runner?.rdf) {
                 return {};
             }
             return this.runner.rdf.outputs[0];
         },
+        additionalParameters() {
+            if (!this?.runner?.rdf) {
+                return [];
+            }
+            const params = this?.runner?.rdf?.additional_parameters || [];
+            console.log("Additional parameters: ", params);
+            return params
+        },
     },
     watch: {
         tileSizes: {
@@ -193,7 +209,17 @@ export default {
                 console.log(oldObj, newObj);
             },
             deep: true
-        }
+        },
+        additionalParameters() {
+            const initValues = {};
+            for (const paramGroup of this.additionalParameters) {
+                for (const param of paramGroup.parameters) {
+                    initValues[param.name] = param.default;
+                }
+            }
+            console.log("Init additional parameters", initValues)
+            this.additionalParameterValues = initValues;
+        },
     },
     mounted() {
         this.init();
@@ -221,11 +247,11 @@ export default {
             this.turnButtons(true);
         },
         async initModel(modelId, runner=undefined) {
-            this.setInfoPanel(`Initializing model ${modelId}...`, true);
-            this.turnButtons(false);
             if (runner === undefined) {
                 runner = this.runner;
             }
+            this.setInfoPanel(`Initializing model ${modelId}...`, true);
+            this.turnButtons(false);
             try {
                 await runner.loadModel(modelId);
             } catch (e) {
@@ -239,10 +265,11 @@ export default {
             this.setInfoPanel("");
         },
         turnButtons(on) {
-            if (!this.runner.rdf) {
+            if (!this?.runner?.rdf) {
                 this.buttonEnabledRun = false;
                 this.buttonEnabledInput = false;
                 this.buttonEnabledOutput = false;
+                this.modelSelectEnable = false;
             } else {
                 this.buttonEnabledRun = on;
                 this.buttonEnabledInput = on && (
@@ -251,10 +278,13 @@ export default {
                 this.buttonEnabledOutput = on && (
                     rdfHas(this.runner.rdf, "test_outputs") ||
                     rdfHas(this.runner.rdf, "sample_outputs"));
+                this.modelSelectEnable = on;
             }
         },
         handleModelChange(model) {
-            this.initModel(model.id);
+            if (this?.runner?.rdf) {
+                this.initModel(model.id);
+            }
         },
         async handleServerUrlChange(url) {
             this.serverUrl = url;
@@ -327,9 +357,10 @@ export default {
                     this.outputSpec,
                     this.tileSizes,
                     this.tileOverlaps,
+                    this.additionalParameterValues,
                     (msg) => {
                         this.setInfoPanel(msg, true);
-                    }
+                    },
                 );
                 if (this.runner.isImg2Img()) {
                     const imgsForShow = processForShow(outTensor, this.outputSpec.axes);
@@ -413,6 +444,6 @@ export default {
             this.setInfoPanel("");
         }
     },
-    components: { HideContainer, OverlayContainer, AdvanceSetting, LoadingAnimation, ModelSelect }
+    components: { HideContainer, OverlayContainer, AdvanceSetting, LoadingAnimation, ModelSelect, ModelParameters }
 };
 </script>
diff --git a/src/modelRun.js b/src/modelRun.js
@@ -59,9 +59,10 @@ class TritonExecutor {
     return ret;
   }
 
-  async runCellpose(array, diameter = 30) {
+  async runCellpose(array, additionalParameters = {}) {
+    console.log("Running cellpose with parameters: ", additionalParameters);
     const ret = await this.triton.execute({
-      inputs: [array, { diameter: diameter }],
+      inputs: [array, additionalParameters],
       model_name: "cellpose-python",
       decode_json: true,
       _rkwargs: true,
@@ -228,18 +229,41 @@ export class ModelRunner {
       sample_inputs: [
         "https://zenodo.org/api/records/6647674/files/sample_input_0.tif/content",
       ],
+      additional_parameters: [
+        {
+          name: "Cellpose Parameters",
+          parameters: [
+            {
+              name: "diameter",
+              type: "number",
+              default: 30,
+              description: "Diameter of the nuclei in pixels.",
+            },
+            {
+              name: "model_type",
+              type: "string",
+              default: "nuclei",
+              description: "Type of cells to segment.",
+              enum: ["nuclei", "cyto"],
+            },
+          ],
+        },
+      ],
     };
     return rdf;
   }
 
-  async submitTensor(tensor) {
+  async submitTensor(tensor, additionalParameters = undefined) {
     const reverseEnd = this.inputEndianness === "<";
     const data_type = this.rdf.inputs[0].data_type;
     const reshapedImg = tfjsToImJoy(tensor, reverseEnd, data_type);
     const modelId = this.modelId;
     let outImg;
     if (modelId === "cellpose-python") {
-      const resp = await this.tritonExecutor.runCellpose(reshapedImg);
+      const resp = await this.tritonExecutor.runCellpose(
+        reshapedImg,
+        additionalParameters
+      );
       outImg = resp.mask;
     } else {
       const resp = await this.tritonExecutor.execute(modelId, [reshapedImg]);
@@ -251,11 +275,11 @@ export class ModelRunner {
     return outImg;
   }
 
-  async runOneTensor(tensor, padder) {
+  async runOneTensor(tensor, padder, additionalParameters = undefined) {
     console.log("Input tile shape: " + tensor.shape);
     const [paddedTensor, padArr] = padder.pad(tensor);
     console.log("Padded tile shape: " + paddedTensor.shape);
-    let outImg = await this.submitTensor(paddedTensor);
+    let outImg = await this.submitTensor(paddedTensor, additionalParameters);
     console.log("Output tile shape: " + outImg._rshape);
     const outTensor = imjoyToTfjs(outImg);
     const isImg2Img =
@@ -275,6 +299,7 @@ export class ModelRunner {
     outputSpec,
     tileSizes,
     tileOverlaps,
+    additionalParameters = undefined,
     reportFunc = undefined
   ) {
     if (!reportFunc) {
@@ -308,7 +333,11 @@ export class ModelRunner {
       const tile = inTiles[i];
       console.log(tile);
       tile.slice(tensor);
-      const outTensor = await this.runOneTensor(tile.data, padder);
+      const outTensor = await this.runOneTensor(
+        tile.data,
+        padder,
+        additionalParameters
+      );
       const outTile = new ImgTile(tile.starts, tile.ends, tile.indexes);
       outTile.data = outTensor;
       outTiles.push(outTile);