Enable custom AI Instance Segmentation jobs (#8849)

This PR does three things:

1. Adds support for training and running AI instance segmentation
models. This is already supported by the worker by was not available
from the UI yet. This workflow is distinct from the regular neuron model
training is designed for nuclei and other instance segmentation tasks.
There is also a lot of naming confusion between our pre-trained models,
the custom model that one can train, generic instance segmentation and
EM_nuclei models (essentially the same) etc. Something for a follow up
PR.
2. Adds AI job settings to better support VX instance segmentations (See
issue #8278):
  - Instance Model Training: `max_distance_nm` 
  - Instance Model Inference: `seed_generator_distance_threshold`
3. It splits the monolithic `starts_jobs_modal.tsx` into separate files
for components, tabs, forms, hooks etc.
  - most of the code is simply moved into sub components
  - added React.callbacks to some callbacks and click handlers


### URL of deployed dev instance (used for testing):
- https://___.webknossos.xyz

### Steps to test:
- Enable worker

Training a new Instances Model:
1. Create annotation with bounding box, open "AI" modal
2. Switch to tab "Train a Model" to start a new model training. Switch
"Model category" to "EM Instance Segmentation for Nuclei, ..."
3. Enter a distance threshold - or use default value
5. Start training

Using an instances model:
1. Create annotation, open "AI" modal
2. Switch to tab "Run a Model"
3. Switch from pre-trained to custom models with toggle switch in upper
right
4. Select previously trained instance model from dropdown
6. Confirm that the "seed generatoion distance" option is shown
7. Start the inference

### Issues:
- fixes #8278

------
(Please delete unneeded items, merge only when none are left open)
- [x] Added changelog entry (create a `$PR_NUMBER.md` file in
`unreleased_changes` or use `./tools/create-changelog-entry.py`)
- [ ] Added migration guide entry if applicable (edit the same file as
for the changelog)
- [ ] Updated [documentation](../blob/master/docs) if applicable
- [ ] Adapted [wk-libs python
client](https://github.com/scalableminds/webknossos-libs/tree/master/webknossos/webknossos/client)
if relevant API parts change
- [ ] Removed dev-only changes like prints and application.conf edits
- [x] Considered [common edge
cases](../blob/master/.github/common_edge_cases.md)
- [ ] Needs datastore update after deployment

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: MichaelBuessemeyer <[email protected]>

Author: 3 people

app/controllers/AiModelController.scala

@@ -30,14 +30,25 @@ object TrainingAnnotationSpecification {
   implicit val jsonFormat: OFormat[TrainingAnnotationSpecification] = Json.format[TrainingAnnotationSpecification]
 }
 
-case class RunTrainingParameters(trainingAnnotations: List[TrainingAnnotationSpecification],
-                                 name: String,
-                                 comment: Option[String],
-                                 aiModelCategory: Option[AiModelCategory],
-                                 workflowYaml: Option[String])
+case class RunNeuronModelTrainingParameters(trainingAnnotations: List[TrainingAnnotationSpecification],
+                                            name: String,
+                                            aiModelCategory: Option[AiModelCategory],
+                                            comment: Option[String],
+                                            workflowYaml: Option[String])
 
-object RunTrainingParameters {
-  implicit val jsonFormat: OFormat[RunTrainingParameters] = Json.format[RunTrainingParameters]
+object RunNeuronModelTrainingParameters {
+  implicit val jsonFormat: OFormat[RunNeuronModelTrainingParameters] = Json.format[RunNeuronModelTrainingParameters]
+}
+
+case class RunInstanceModelTrainingParameters(trainingAnnotations: List[TrainingAnnotationSpecification],
+                                              name: String,
+                                              aiModelCategory: Option[AiModelCategory],
+                                              maxDistanceNm: Option[Double],
+                                              comment: Option[String],
+                                              workflowYaml: Option[String])
+
+object RunInstanceModelTrainingParameters {
+  implicit val jsonFormat: OFormat[RunInstanceModelTrainingParameters] = Json.format[RunInstanceModelTrainingParameters]
 }
 
 case class RunInferenceParameters(annotationId: Option[ObjectId],
@@ -48,7 +59,8 @@ case class RunInferenceParameters(annotationId: Option[ObjectId],
                                   boundingBox: String,
                                   newDatasetName: String,
                                   maskAnnotationLayerName: Option[String],
-                                  workflowYaml: Option[String])
+                                  workflowYaml: Option[String],
+                                  seedGeneratorDistanceThreshold: Option[Double])
 
 object RunInferenceParameters {
   implicit val jsonFormat: OFormat[RunInferenceParameters] = Json.format[RunInferenceParameters]
@@ -127,21 +139,18 @@ class AiModelController @Inject()(
     }
   }
 
-  def runNeuronTraining: Action[RunTrainingParameters] = sil.SecuredAction.async(validateJson[RunTrainingParameters]) {
-    implicit request =>
+  def runNeuronTraining: Action[RunNeuronModelTrainingParameters] =
+    sil.SecuredAction.async(validateJson[RunNeuronModelTrainingParameters]) { implicit request =>
       for {
         _ <- userService.assertIsSuperUser(request.identity)
         trainingAnnotations = request.body.trainingAnnotations
-        _ <- Fox
-          .fromBool(trainingAnnotations.nonEmpty || request.body.workflowYaml.isDefined) ?~> "aiModel.training.zeroAnnotations"
+        _ <- Fox.fromBool(trainingAnnotations.nonEmpty || request.body.workflowYaml.isDefined) ?~> "aiModel.training.zeroAnnotations"
         firstAnnotationId <- trainingAnnotations.headOption.map(_.annotationId).toFox
         annotation <- annotationDAO.findOne(firstAnnotationId)
         dataset <- datasetDAO.findOne(annotation._dataset)
-        _ <- Fox
-          .fromBool(request.identity._organization == dataset._organization) ?~> "job.trainModel.notAllowed.organization" ~> FORBIDDEN
+        _ <- Fox.fromBool(request.identity._organization == dataset._organization) ?~> "job.trainModel.notAllowed.organization" ~> FORBIDDEN
         dataStore <- dataStoreDAO.findOneByName(dataset._dataStore) ?~> "dataStore.notFound"
-        _ <- Fox
-          .serialCombined(request.body.trainingAnnotations.map(_.annotationId))(annotationDAO.findOne) ?~> "annotation.notFound"
+        _ <- Fox.serialCombined(request.body.trainingAnnotations.map(_.annotationId))(annotationDAO.findOne) ?~> "annotation.notFound"
         modelId = ObjectId.generate
         organization <- organizationDAO.findOne(request.identity._organization)
         jobCommand = JobCommand.train_neuron_model
@@ -154,8 +163,7 @@ class AiModelController @Inject()(
         existingAiModelsCount <- aiModelDAO.countByNameAndOrganization(request.body.name,
                                                                        request.identity._organization)
         _ <- Fox.fromBool(existingAiModelsCount == 0) ?~> "aiModel.nameInUse"
-        newTrainingJob <- jobService
-          .submitJob(jobCommand, commandArgs, request.identity, dataStore.name) ?~> "job.couldNotRunTrainModel"
+        newTrainingJob <- jobService.submitJob(jobCommand, commandArgs, request.identity, dataStore.name) ?~> "job.couldNotRunTrainModel"
         newAiModel = AiModel(
           _id = modelId,
           _organization = request.identity._organization,
@@ -171,7 +179,93 @@ class AiModelController @Inject()(
         _ <- aiModelDAO.insertOne(newAiModel)
         newAiModelJs <- aiModelService.publicWrites(newAiModel, request.identity)
       } yield Ok(newAiModelJs)
-  }
+    }
+
+  def runInstanceTraining: Action[RunInstanceModelTrainingParameters] =
+    sil.SecuredAction.async(validateJson[RunInstanceModelTrainingParameters]) { implicit request =>
+      for {
+        _ <- userService.assertIsSuperUser(request.identity)
+        trainingAnnotations = request.body.trainingAnnotations
+        _ <- Fox.fromBool(trainingAnnotations.nonEmpty || request.body.workflowYaml.isDefined) ?~> "aiModel.training.zeroAnnotations"
+        firstAnnotationId <- trainingAnnotations.headOption.map(_.annotationId).toFox
+        annotation <- annotationDAO.findOne(firstAnnotationId)
+        dataset <- datasetDAO.findOne(annotation._dataset)
+        _ <- Fox.fromBool(request.identity._organization == dataset._organization) ?~> "job.trainModel.notAllowed.organization" ~> FORBIDDEN
+        dataStore <- dataStoreDAO.findOneByName(dataset._dataStore) ?~> "dataStore.notFound"
+        _ <- Fox.serialCombined(request.body.trainingAnnotations.map(_.annotationId))(annotationDAO.findOne) ?~> "annotation.notFound"
+        modelId = ObjectId.generate
+        organization <- organizationDAO.findOne(request.identity._organization)
+        jobCommand = JobCommand.train_instance_model
+        commandArgs = Json.obj(
+          "training_annotations" -> Json.toJson(trainingAnnotations),
+          "organization_id" -> organization._id,
+          "model_id" -> modelId,
+          "custom_workflow_provided_by_user" -> request.body.workflowYaml,
+          "max_distance_nm" -> request.body.maxDistanceNm
+        )
+        existingAiModelsCount <- aiModelDAO.countByNameAndOrganization(request.body.name,
+                                                                       request.identity._organization)
+        _ <- Fox.fromBool(existingAiModelsCount == 0) ?~> "aiModel.nameInUse"
+        newTrainingJob <- jobService.submitJob(jobCommand, commandArgs, request.identity, dataStore.name) ?~> "job.couldNotRunTrainModel"
+        newAiModel = AiModel(
+          _id = modelId,
+          _organization = request.identity._organization,
+          _sharedOrganizations = List(),
+          _dataStore = dataStore.name,
+          _user = request.identity._id,
+          _trainingJob = Some(newTrainingJob._id),
+          _trainingAnnotations = trainingAnnotations.map(_.annotationId),
+          name = request.body.name,
+          comment = request.body.comment,
+          category = request.body.aiModelCategory
+        )
+        _ <- aiModelDAO.insertOne(newAiModel)
+        newAiModelJs <- aiModelService.publicWrites(newAiModel, request.identity)
+      } yield Ok(newAiModelJs)
+    }
+
+  def runCustomInstanceModelInference: Action[RunInferenceParameters] =
+    sil.SecuredAction.async(validateJson[RunInferenceParameters]) { implicit request =>
+      for {
+        _ <- userService.assertIsSuperUser(request.identity)
+        organization <- organizationDAO.findOne(request.body.organizationId)(GlobalAccessContext) ?~> Messages(
+          "organization.notFound",
+          request.body.organizationId)
+        _ <- Fox.fromBool(request.identity._organization == organization._id) ?~> "job.runInference.notAllowed.organization" ~> FORBIDDEN
+        dataset <- datasetDAO.findOneByDirectoryNameAndOrganization(request.body.datasetDirectoryName, organization._id)
+        dataStore <- dataStoreDAO.findOneByName(dataset._dataStore) ?~> "dataStore.notFound"
+        _ <- aiModelDAO.findOne(request.body.aiModelId) ?~> "aiModel.notFound"
+        _ <- datasetService.assertValidDatasetName(request.body.newDatasetName)
+        jobCommand = JobCommand.infer_instances
+        boundingBox <- BoundingBox.fromLiteral(request.body.boundingBox).toFox
+        commandArgs = Json.obj(
+          "dataset_id" -> dataset._id,
+          "organization_id" -> organization._id,
+          "dataset_name" -> dataset.name,
+          "layer_name" -> request.body.colorLayerName,
+          "bbox" -> boundingBox.toLiteral,
+          "model_id" -> request.body.aiModelId,
+          "dataset_directory_name" -> request.body.datasetDirectoryName,
+          "new_dataset_name" -> request.body.newDatasetName,
+          "custom_workflow_provided_by_user" -> request.body.workflowYaml,
+          "seed_generator_distance_threshold" -> request.body.seedGeneratorDistanceThreshold
+        )
+        newInferenceJob <- jobService.submitJob(jobCommand, commandArgs, request.identity, dataStore.name) ?~> "job.couldNotRunInferWithModel"
+        newAiInference = AiInference(
+          _id = ObjectId.generate,
+          _organization = request.identity._organization,
+          _aiModel = request.body.aiModelId,
+          _newDataset = None,
+          _annotation = request.body.annotationId,
+          boundingBox = boundingBox,
+          _inferenceJob = newInferenceJob._id,
+          newSegmentationLayerName = "segmentation",
+          maskAnnotationLayerName = request.body.maskAnnotationLayerName
+        )
+        _ <- aiInferenceDAO.insertOne(newAiInference)
+        newAiModelJs <- aiInferenceService.publicWrites(newAiInference, request.identity)
+      } yield Ok(newAiModelJs)
+    }
 
   def runCustomNeuronInference: Action[RunInferenceParameters] =
     sil.SecuredAction.async(validateJson[RunInferenceParameters]) { implicit request =>

app/models/aimodels/AiModel.scala

@@ -71,7 +71,8 @@ class AiModelService @Inject()(dataStoreDAO: DataStoreDAO,
         "comment" -> aiModel.comment,
         "trainingJob" -> trainingJobJsOpt,
         "created" -> aiModel.created,
-        "sharedOrganizationIds" -> sharedOrganizationIds
+        "sharedOrganizationIds" -> sharedOrganizationIds,
+        "category" -> aiModel.category
       )
 }
 

app/models/job/JobCommand.scala

@@ -12,8 +12,8 @@ object JobCommand extends ExtendedEnumeration {
    */
 
   val compute_mesh_file, compute_segment_index_file, convert_to_wkw, export_tiff, find_largest_segment_id,
-  globalize_floodfills, infer_nuclei, infer_neurons, materialize_volume_annotation, render_animation,
-  infer_mitochondria, align_sections, train_model, infer_with_model, train_neuron_model = Value
+  globalize_floodfills, infer_nuclei, infer_neurons, infer_instances, materialize_volume_annotation, render_animation,
+  infer_mitochondria, align_sections, train_model, infer_with_model, train_neuron_model, train_instance_model = Value
 
   val highPriorityJobs: Set[Value] = Set(convert_to_wkw, export_tiff)
   val lowPriorityJobs: Set[Value] = values.diff(highPriorityJobs)

conf/webknossos.latest.routes

@@ -285,8 +285,10 @@ POST          /jobs/:id/attachDatasetToInference
 GET           /jobs/:id/export                                                                  controllers.JobController.redirectToExport(id: ObjectId)
 
 # AI Models
-POST          /aiModels/runNeuronTraining                                                       controllers.AiModelController.runNeuronTraining
-POST          /aiModels/inferences/runCustomNeuronInference                                     controllers.AiModelController.runCustomNeuronInference
+POST          /aiModels/runNeuronModelTraining                                                  controllers.AiModelController.runNeuronTraining
+POST          /aiModels/runInstanceModelTraining                                                controllers.AiModelController.runInstanceTraining 
+POST          /aiModels/inferences/runCustomNeuronModelInference                                controllers.AiModelController.runCustomNeuronInference
+POST          /aiModels/inferences/runCustomInstanceModelInference                              controllers.AiModelController.runCustomInstanceModelInference
 GET           /aiModels/inferences/:id                                                          controllers.AiModelController.readAiInferenceInfo(id: ObjectId)
 GET           /aiModels/inferences                                                              controllers.AiModelController.listAiInferences
 GET           /aiModels                                                                         controllers.AiModelController.listAiModels

frontend/javascripts/admin/api/jobs.ts

@@ -351,7 +351,11 @@ export function startAlignSectionsJob(
   });
 }
 
-type AiModelCategory = "em_neurons" | "em_nuclei";
+// This enum needs to be kept in sync with the backend/database
+export enum APIAiModelCategory {
+  EM_NEURONS = "em_neurons",
+  EM_NUCLEI = "em_nuclei",
+}
 
 type AiModelTrainingAnnotationSpecification = {
   annotationId: string;
@@ -360,22 +364,38 @@ type AiModelTrainingAnnotationSpecification = {
   mag: Vector3;
 };
 
-type RunTrainingParameters = {
-  trainingAnnotations: Array<AiModelTrainingAnnotationSpecification>;
+type RunNeuronModelTrainingParameters = {
+  trainingAnnotations: AiModelTrainingAnnotationSpecification[];
+  name: string;
+  aiModelCategory: APIAiModelCategory.EM_NEURONS;
+  comment?: string;
+  workflowYaml?: string;
+};
+
+export function runNeuronTraining(params: RunNeuronModelTrainingParameters) {
+  return Request.sendJSONReceiveJSON("/api/aiModels/runNeuronModelTraining", {
+    method: "POST",
+    data: JSON.stringify(params),
+  });
+}
+
+type RunInstanceModelTrainingParameters = {
+  trainingAnnotations: AiModelTrainingAnnotationSpecification[];
   name: string;
+  aiModelCategory: APIAiModelCategory.EM_NUCLEI;
+  maxDistanceNm: number;
   comment?: string;
-  aiModelCategory?: AiModelCategory;
   workflowYaml?: string;
 };
 
-export function runNeuronTraining(params: RunTrainingParameters) {
-  return Request.sendJSONReceiveJSON("/api/aiModels/runNeuronTraining", {
+export function runInstanceModelTraining(params: RunInstanceModelTrainingParameters) {
+  return Request.sendJSONReceiveJSON("/api/aiModels/runInstanceModelTraining", {
     method: "POST",
     data: JSON.stringify(params),
   });
 }
 
-type RunInferenceParameters = {
+export type BaseModelInferenceParameters = {
   annotationId?: string;
   aiModelId: string;
   datasetDirectoryName: string;
@@ -386,9 +406,23 @@ type RunInferenceParameters = {
   workflowYaml?: string;
   // maskAnnotationLayerName?: string | null
 };
+type RunNeuronModelInferenceParameters = BaseModelInferenceParameters;
+
+type RunInstanceModelInferenceParameters = BaseModelInferenceParameters & {
+  seedGeneratorDistanceThreshold: number;
+};
+
+export function runNeuronModelInferenceWithAiModelJob(params: RunNeuronModelInferenceParameters) {
+  return Request.sendJSONReceiveJSON("/api/aiModels/inferences/runCustomNeuronModelInference", {
+    method: "POST",
+    data: JSON.stringify({ ...params, boundingBox: params.boundingBox.join(",") }),
+  });
+}
 
-export function runNeuronInferenceWithAiModelJob(params: RunInferenceParameters) {
-  return Request.sendJSONReceiveJSON("/api/aiModels/inferences/runCustomNeuronInference", {
+export function runInstanceModelInferenceWithAiModelJob(
+  params: RunInstanceModelInferenceParameters,
+) {
+  return Request.sendJSONReceiveJSON("/api/aiModels/inferences/runCustomInstanceModelInference", {
     method: "POST",
     data: JSON.stringify({ ...params, boundingBox: params.boundingBox.join(",") }),
   });

frontend/javascripts/admin/job/job_list_view.tsx

@@ -126,7 +126,7 @@ export function JobState({ job }: { job: APIJob }) {
 
   return (
     <Tooltip title={tooltip}>
-      <span className="icon-margin-right">{icon}</span>
+      <span>{icon}</span>
       {jobStateNormalized}
     </Tooltip>
   );
@@ -235,7 +235,7 @@ function JobListView() {
     ) {
       return (
         <span>
-          Neuron inferral for layer {job.layerName} of{" "}
+          AI Neuron inferral for layer <i>{job.layerName}</i> of{" "}
           <Link to={linkToDataset}>{job.datasetName}</Link>{" "}
         </span>
       );
@@ -256,14 +256,21 @@ function JobListView() {
     ) {
       return (
         <span>
-          Mitochondria inferral for layer {job.layerName} of{" "}
+          AI Mitochondria inferral for layer <i>{job.layerName}</i> of{" "}
+          <Link to={linkToDataset}>{job.datasetName}</Link>{" "}
+        </span>
+      );
+    } else if (job.type === APIJobType.INFER_INSTANCES && linkToDataset != null && job.layerName) {
+      return (
+        <span>
+          AI instance segmentation for layer <i>{job.layerName}</i> of{" "}
           <Link to={linkToDataset}>{job.datasetName}</Link>{" "}
         </span>
       );
     } else if (job.type === APIJobType.ALIGN_SECTIONS && linkToDataset != null && job.layerName) {
       return (
         <span>
-          Align sections for layer {job.layerName} of{" "}
+          Align sections for layer <i>{job.layerName}</i> of{" "}
           <Link to={linkToDataset}>{job.datasetName}</Link>{" "}
         </span>
       );
@@ -277,11 +284,19 @@ function JobListView() {
             : null}
         </span>
       );
-    } else if (job.type === APIJobType.TRAIN_NEURON_MODEL || APIJobType.DEPRECATED_TRAIN_MODEL) {
-      const numberOfTrainingAnnotations = job.trainingAnnotations.length;
+    } else if (
+      job.type === APIJobType.TRAIN_NEURON_MODEL ||
+      job.type === APIJobType.TRAIN_INSTANCE_MODEL ||
+      job.type === APIJobType.DEPRECATED_TRAIN_MODEL
+    ) {
+      const numberOfTrainingAnnotations = job.trainingAnnotations?.length || 0;
+      const modelName =
+        job.type === APIJobType.TRAIN_NEURON_MODEL || job.type === APIJobType.DEPRECATED_TRAIN_MODEL
+          ? "neuron model"
+          : "instance model";
       return (
         <span>
-          {`Train neuron model on ${numberOfTrainingAnnotations} ${Utils.pluralize("annotation", numberOfTrainingAnnotations)}. `}
+          {`Train ${modelName} on ${numberOfTrainingAnnotations} ${Utils.pluralize("annotation", numberOfTrainingAnnotations)}. `}
           {getShowTrainingDataLink(job.trainingAnnotations)}
         </span>
       );