Enable custom AI Instance Segmentation jobs

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>
       );