Evaluate segmentation in infer neurons task (#8221)

* added new inputs for model evaluation * run formatter * updated changelog * add support for displaying annotation links instead of datasets * Update CHANGELOG.unreleased.md Co-authored-by: MichaelBuessemeyer <[email protected]> * adapt to feedback * update antd to allow form.field to set its own layout * Update frontend/javascripts/admin/api/jobs.ts Co-authored-by: MichaelBuessemeyer <[email protected]> * apply some of the requested changes * yarn fix-frontend * add assertion for annotationId * change variable name in error message * remove changes in application.conf * actually undo changes in application.conf * format frontend * apply feedback * fix drag handle alignment --------- Co-authored-by: MichaelBuessemeyer <[email protected]> Co-authored-by: Michael Büßemeyer <[email protected]>
scalableminds · Jan 31, 2025 · a723fcd · a723fcd
1 parent bd70862
commit a723fcd
Show file tree

Hide file tree

Showing 10 changed files with 285 additions and 71 deletions.
diff --git a/CHANGELOG.unreleased.md b/CHANGELOG.unreleased.md
@@ -11,6 +11,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 [Commits](https://github.com/scalableminds/webknossos/compare/25.01.0...HEAD)
 
 ### Added
+- It is now possible to start a split-merger evaluation when starting a neuron inference. [#8221](https://github.com/scalableminds/webknossos/pull/8221)
 - Added the possibility to configure a rotation for a dataset, which can be toggled off and on when viewing and annotating data. [#8159](https://github.com/scalableminds/webknossos/pull/8159)
 - When using the “Restore older Version” feature, there are no longer separate tabs for the different annotation layers. Only one linear annotation history is now used, and if you revert to an older version, all layers are reverted. If layers were added/deleted since then, that is also reverted. This also means that proofreading annotations can now be reverted to older versions as well. The description text of annotations is now versioned as well. [#7917](https://github.com/scalableminds/webknossos/pull/7917)
 - Added the possibility to use the "merger mode" even when the user has annotated volume data in the current layer (as long as no other mapping is active). [#8335](https://github.com/scalableminds/webknossos/pull/8335)

diff --git a/app/controllers/JobController.scala b/app/controllers/JobController.scala
@@ -223,10 +223,17 @@ class JobController @Inject()(
   def runInferNeuronsJob(datasetId: String,
                          layerName: String,
                          bbox: String,
-                         newDatasetName: String): Action[AnyContent] =
+                         newDatasetName: String,
+                         doSplitMergerEvaluation: Boolean,
+                         annotationId: Option[String],
+                         evalUseSparseTracing: Option[Boolean],
+                         evalMaxEdgeLength: Option[Double],
+                         evalSparseTubeThresholdNm: Option[Double],
+                         evalMinMergerPathLengthNm: Option[Double]): Action[AnyContent] =
     sil.SecuredAction.async { implicit request =>
       log(Some(slackNotificationService.noticeFailedJobRequest)) {
         for {
+
           datasetIdValidated <- ObjectId.fromString(datasetId)
           dataset <- datasetDAO.findOne(datasetIdValidated) ?~> Messages("dataset.notFound", datasetId) ~> NOT_FOUND
           organization <- organizationDAO.findOne(dataset._organization)(GlobalAccessContext) ?~> Messages(
@@ -237,6 +244,7 @@ class JobController @Inject()(
           _ <- datasetService.assertValidLayerNameLax(layerName)
           multiUser <- multiUserDAO.findOne(request.identity._multiUser)
           _ <- Fox.runIf(!multiUser.isSuperUser)(jobService.assertBoundingBoxLimits(bbox, None))
+          annotationIdParsed <- Fox.runIf(doSplitMergerEvaluation)(annotationId.toFox) ?~> "job.inferNeurons.annotationIdEvalParamsMissing"
           command = JobCommand.infer_neurons
           commandArgs = Json.obj(
             "organization_id" -> organization._id,
@@ -245,6 +253,12 @@ class JobController @Inject()(
             "new_dataset_name" -> newDatasetName,
             "layer_name" -> layerName,
             "bbox" -> bbox,
+            "do_split_merger_evaluation" -> doSplitMergerEvaluation,
+            "annotation_id" -> annotationIdParsed,
+            "eval_use_sparse_tracing" -> evalUseSparseTracing,
+            "eval_max_edge_length" -> evalMaxEdgeLength,
+            "eval_sparse_tube_threshold_nm" -> evalSparseTubeThresholdNm,
+            "eval_min_merger_path_length_nm" -> evalMinMergerPathLengthNm,
           )
           job <- jobService.submitJob(command, commandArgs, request.identity, dataset._dataStore) ?~> "job.couldNotRunNeuronInferral"
           js <- jobService.publicWrites(job)

diff --git a/app/models/job/Job.scala b/app/models/job/Job.scala
@@ -54,6 +54,7 @@ case class Job(
   def datasetId: Option[String] = argAsStringOpt("dataset_id")
 
   private def argAsStringOpt(key: String) = (commandArgs \ key).toOption.flatMap(_.asOpt[String])
+  private def argAsBooleanOpt(key: String) = (commandArgs \ key).toOption.flatMap(_.asOpt[Boolean])
 
   def resultLink(organizationId: String): Option[String] =
     if (effectiveState != JobState.SUCCESS) None
@@ -66,6 +67,10 @@ case class Job(
           }.getOrElse(datasetName.map(name => s"datasets/$organizationId/$name/view"))
         case JobCommand.export_tiff | JobCommand.render_animation =>
           Some(s"/api/jobs/${this._id}/export")
+        case JobCommand.infer_neurons if this.argAsBooleanOpt("do_evaluation").getOrElse(false) =>
+          returnValue.map { resultAnnotationLink =>
+            resultAnnotationLink
+          }
         case JobCommand.infer_nuclei | JobCommand.infer_neurons | JobCommand.materialize_volume_annotation |
             JobCommand.infer_with_model | JobCommand.infer_mitochondria | JobCommand.align_sections =>
           // Old jobs before the dataset renaming changes returned the output dataset name.

diff --git a/conf/messages b/conf/messages
@@ -331,6 +331,7 @@ job.edgeLengthExceeded = An edge length of the selected bounding box is too larg
 job.convertToWkw.notAllowed.organization = Converting to WKW is only allowed for datasets of your own organization.
 job.inferNuclei.notAllowed.organization = Currently nuclei inferral is only allowed for datasets of your own organization.
 job.inferNeurons.notAllowed.organization = Currently neuron inferral is only allowed for datasets of your own organization.
+job.inferNeurons.annotationIdEvalParamsMissing=A evaluation of the neuron inferral jobs was requested but no annotation was supplied.
 job.meshFile.notAllowed.organization = Calculating mesh files is only allowed for datasets of your own organization.
 job.segmentIndexFile.notAllowed.organization = Calculating segment index files is only allowed for datasets of your own organization.
 job.globalizeFloodfill.notAllowed.organization = Globalizing floodfills is only allowed for datasets of your own organization.

diff --git a/conf/webknossos.latest.routes b/conf/webknossos.latest.routes
@@ -260,7 +260,7 @@ POST          /jobs/run/computeMeshFile/:datasetId
 POST          /jobs/run/computeSegmentIndexFile/:datasetId                                      controllers.JobController.runComputeSegmentIndexFileJob(datasetId: String, layerName: String)
 POST          /jobs/run/exportTiff/:datasetId                                                   controllers.JobController.runExportTiffJob(datasetId: String, bbox: String, additionalCoordinates: Option[String], layerName: Option[String], mag: Option[String], annotationLayerName: Option[String], annotationId: Option[String], asOmeTiff: Boolean)
 POST          /jobs/run/inferNuclei/:datasetId                                                  controllers.JobController.runInferNucleiJob(datasetId: String, layerName: String, newDatasetName: String)
-POST          /jobs/run/inferNeurons/:datasetId                                                 controllers.JobController.runInferNeuronsJob(datasetId: String, layerName: String, bbox: String, newDatasetName: String)
+POST          /jobs/run/inferNeurons/:datasetId                                                 controllers.JobController.runInferNeuronsJob(datasetId: String, layerName: String, bbox: String, newDatasetName: String, doSplitMergerEvaluation: Boolean, annotationId: Option[String], evalUseSparseTracing: Option[Boolean], evalMaxEdgeLength: Option[Double], evalSparseTubeThresholdNm: Option[Double], evalMinMergerPathLengthNm: Option[Double])
 POST          /jobs/run/inferMitochondria/:datasetId                                            controllers.JobController.runInferMitochondriaJob(datasetId: String, layerName: String, bbox: String, newDatasetName: String)
 POST          /jobs/run/alignSections/:datasetId                                                controllers.JobController.runAlignSectionsJob(datasetId: String, layerName: String, newDatasetName: String, annotationId: Option[String])
 POST          /jobs/run/materializeVolumeAnnotation/:datasetId                                  controllers.JobController.runMaterializeVolumeAnnotationJob(datasetId: String, fallbackLayerName: String, annotationId: String, annotationType: String, newDatasetName: String, outputSegmentationLayerName: String, mergeSegments: Boolean, volumeLayerName: Option[String], includesEditableMapping: Boolean, boundingBox: Option[String])

diff --git a/frontend/javascripts/admin/api/jobs.ts b/frontend/javascripts/admin/api/jobs.ts
@@ -179,12 +179,37 @@ export function startNeuronInferralJob(
   layerName: string,
   bbox: Vector6,
   newDatasetName: string,
+  doSplitMergerEvaluation: boolean,
+  annotationId?: string,
+  useSparseTracing?: boolean,
+  evalMaxEdgeLength?: number,
+  evalSparseTubeThresholdNm?: number,
+  evalMinMergerPathLengthNm?: number,
 ): Promise<APIJob> {
   const urlParams = new URLSearchParams({
     layerName,
     bbox: bbox.join(","),
     newDatasetName,
+    doSplitMergerEvaluation: doSplitMergerEvaluation.toString(),
   });
+  if (doSplitMergerEvaluation) {
+    if (!annotationId) {
+      throw new Error("annotationId is required when doSplitMergerEvaluation is true");
+    }
+    urlParams.append("annotationId", `${annotationId}`);
+    if (useSparseTracing != null) {
+      urlParams.append("evalUseSparseTracing", `${useSparseTracing}`);
+    }
+    if (evalMaxEdgeLength != null) {
+      urlParams.append("evalMaxEdgeLength", `${evalMaxEdgeLength}`);
+    }
+    if (evalSparseTubeThresholdNm != null) {
+      urlParams.append("evalSparseTubeThresholdNm", `${evalSparseTubeThresholdNm}`);
+    }
+    if (evalMinMergerPathLengthNm != null) {
+      urlParams.append("evalMinMergerPathLengthNm", `${evalMinMergerPathLengthNm}`);
+    }
+  }
   return Request.receiveJSON(`/api/jobs/run/inferNeurons/${datasetId}?${urlParams.toString()}`, {
     method: "POST",
   });

diff --git a/frontend/javascripts/oxalis/view/action-bar/starting_job_modals.tsx b/frontend/javascripts/oxalis/view/action-bar/starting_job_modals.tsx
@@ -14,8 +14,11 @@ import {
   Button,
   Card,
   Checkbox,
+  Col,
+  Collapse,
   Form,
   type FormInstance,
+  InputNumber,
   Modal,
   Radio,
   Row,
@@ -98,6 +101,7 @@ type StartJobFormProps = Props & {
   jobApiCall: (arg0: JobApiCallArgsType, form: FormInstance<any>) => Promise<void | APIJob>;
   jobName: keyof typeof jobNameToImagePath;
   description: React.ReactNode;
+  jobSpecificInputFields?: React.ReactNode | undefined;
   isBoundingBoxConfigurable?: boolean;
   chooseSegmentationLayer?: boolean;
   suggestedDatasetSuffix: string;
@@ -523,11 +527,92 @@ function ShouldUseTreesFormItem() {
   );
 }
 
+type SplitMergerEvaluationSettings = {
+  useSparseTracing?: boolean;
+  maxEdgeLength?: number;
+  sparseTubeThresholdInNm?: number;
+  minimumMergerPathLengthInNm?: number;
+};
+
+function CollapsibleSplitMergerEvaluationSettings({
+  isActive = false,
+  setActive,
+}: { isActive: boolean; setActive: (active: boolean) => void }) {
+  return (
+    <Collapse
+      style={{ marginBottom: 8 }}
+      onChange={() => setActive(!isActive)}
+      expandIcon={() => <Checkbox checked={isActive} />}
+      items={[
+        {
+          key: "evaluation",
+          label: "Evaluation Settings",
+          children: (
+            <Row>
+              <Col style={{ width: "100%" }}>
+                <Form.Item
+                  layout="horizontal"
+                  label="Use sparse ground truth tracing"
+                  name={["splitMergerEvaluationSettings", "useSparseTracing"]}
+                  valuePropName="checked"
+                  initialValue={true}
+                  tooltip="The evaluation mode can either be `dense`
+    in case all processes in the volume are annotated in the ground-truth.
+    If not, use the `sparse` mode."
+                >
+                  <Checkbox style={{ width: "100%" }} />
+                </Form.Item>
+                <Form.Item
+                  label="Max edge length in nm"
+                  name={["splitMergerEvaluationSettings", "maxEdgeLength"]}
+                  tooltip="Ground truth tracings can be densified so that
+    nodes are at most max_edge_length nm apart.
+    However, this can also introduce wrong nodes in curved processes."
+                >
+                  <InputNumber style={{ width: "100%" }} placeholder="None" />
+                </Form.Item>
+                <Form.Item
+                  label="Sparse tube threshold in nm"
+                  name={["splitMergerEvaluationSettings", "sparseTubeThresholdInNm"]}
+                  tooltip="Tube threshold for sparse evaluation,
+    determining if a process is too far from the ground-truth."
+                >
+                  <InputNumber style={{ width: "100%" }} placeholder="1000" />
+                </Form.Item>
+                <Form.Item
+                  label="Sparse minimum merger path length in nm"
+                  name={["splitMergerEvaluationSettings", "minimumMergerPathLengthInNm"]}
+                  tooltip="Minimum ground truth path length of a merger component
+    to be counted as a relevant merger (for sparse evaluation).
+    Note, the path length to neighboring nodes of a component is included for this comparison. This optimistic path length
+    estimation makes sure no relevant mergers are ignored."
+                >
+                  <InputNumber style={{ width: "100%" }} placeholder="800" />
+                </Form.Item>
+                <Form.Item name="useAnnotation" initialValue={true} hidden />
+              </Col>
+            </Row>
+          ),
+        },
+      ]}
+      activeKey={isActive ? "evaluation" : []}
+    />
+  );
+}
+
 function StartJobForm(props: StartJobFormProps) {
   const isBoundingBoxConfigurable = props.isBoundingBoxConfigurable || false;
   const isSkeletonSelectable = props.isSkeletonSelectable || false;
   const chooseSegmentationLayer = props.chooseSegmentationLayer || false;
-  const { handleClose, jobName, jobApiCall, fixedSelectedLayer, title, description } = props;
+  const {
+    handleClose,
+    jobName,
+    jobApiCall,
+    fixedSelectedLayer,
+    title,
+    description,
+    jobSpecificInputFields,
+  } = props;
   const [form] = Form.useForm();
   const rawUserBoundingBoxes = useSelector((state: OxalisState) =>
     getUserBoundingBoxesFromState(state),
@@ -650,6 +735,7 @@ function StartJobForm(props: StartJobFormProps) {
         onChangeSelectedBoundingBox={(bBoxId) => form.setFieldsValue({ boundingBoxId: bBoxId })}
         value={form.getFieldValue("boundingBoxId")}
       />
+      {jobSpecificInputFields}
       {isSkeletonSelectable && <ShouldUseTreesFormItem />}
       {props.showWorkflowYaml ? (
         <CollapsibleWorkflowYamlEditor
@@ -701,7 +787,9 @@ export function NucleiDetectionForm() {
 }
 export function NeuronSegmentationForm() {
   const dataset = useSelector((state: OxalisState) => state.dataset);
+  const hasSkeletonAnnotation = useSelector((state: OxalisState) => state.tracing.skeleton != null);
   const dispatch = useDispatch();
+  const [doSplitMergerEvaluation, setDoSplitMergerEvaluation] = React.useState(false);
   return (
     <StartJobForm
       handleClose={() => dispatch(setAIJobModalStateAction("invisible"))}
@@ -710,13 +798,42 @@ export function NeuronSegmentationForm() {
       title="AI Neuron Segmentation"
       suggestedDatasetSuffix="with_reconstructed_neurons"
       isBoundingBoxConfigurable
-      jobApiCall={async ({ newDatasetName, selectedLayer: colorLayer, selectedBoundingBox }) => {
-        if (!selectedBoundingBox) {
+      jobApiCall={async (
+        { newDatasetName, selectedLayer: colorLayer, selectedBoundingBox, annotationId },
+        form: FormInstance<any>,
+      ) => {
+        const splitMergerEvaluationSettings = form.getFieldValue(
+          "splitMergerEvaluationSettings",
+        ) as SplitMergerEvaluationSettings;
+        if (
+          !selectedBoundingBox ||
+          (doSplitMergerEvaluation && splitMergerEvaluationSettings == null)
+        ) {
           return;
         }
 
         const bbox = computeArrayFromBoundingBox(selectedBoundingBox.boundingBox);
-        return startNeuronInferralJob(dataset.id, colorLayer.name, bbox, newDatasetName);
+        if (!doSplitMergerEvaluation) {
+          return startNeuronInferralJob(
+            dataset.id,
+            colorLayer.name,
+            bbox,
+            newDatasetName,
+            doSplitMergerEvaluation,
+          );
+        }
+        return startNeuronInferralJob(
+          dataset.id,
+          colorLayer.name,
+          bbox,
+          newDatasetName,
+          doSplitMergerEvaluation,
+          annotationId,
+          splitMergerEvaluationSettings.useSparseTracing,
+          splitMergerEvaluationSettings.maxEdgeLength,
+          splitMergerEvaluationSettings.sparseTubeThresholdInNm,
+          splitMergerEvaluationSettings.minimumMergerPathLengthInNm,
+        );
       }}
       description={
         <>
@@ -731,6 +848,14 @@ export function NeuronSegmentationForm() {
           </Space>
         </>
       }
+      jobSpecificInputFields={
+        hasSkeletonAnnotation && (
+          <CollapsibleSplitMergerEvaluationSettings
+            isActive={doSplitMergerEvaluation}
+            setActive={setDoSplitMergerEvaluation}
+          />
+        )
+      }
     />
   );
 }

diff --git a/frontend/javascripts/oxalis/view/left-border-tabs/layer_settings_tab.tsx b/frontend/javascripts/oxalis/view/left-border-tabs/layer_settings_tab.tsx
@@ -184,7 +184,10 @@ function DragHandle({ id }: { id: string }) {
 
 function DummyDragHandle({ tooltipTitle }: { tooltipTitle: string }) {
   return (
-    <FastTooltip title={tooltipTitle}>
+    <FastTooltip
+      title={tooltipTitle}
+      style={{ justifyContent: "center", alignItems: "center", display: "flex" }}
+    >
       <DragHandleIcon isDisabled />
     </FastTooltip>
   );

diff --git a/package.json b/package.json
@@ -141,7 +141,7 @@
     "@tanstack/react-query-persist-client": "4.36.1",
     "@zip.js/zip.js": "^2.7.32",
     "ansi-to-react": "^6.1.6",
-    "antd": "5.17.4",
+    "antd": "5.18",
     "ball-morphology": "^0.1.0",
     "base64-js": "^1.2.1",
     "beautiful-react-hooks": "^3.11.1",