Allow training a model on multiple annotations

frontend/javascripts/admin/api/jobs.ts

@@ -313,7 +313,7 @@ export function startAlignSectionsJob(
 
 type AiModelCategory = "em_neurons" | "em_nuclei";
 
-type AiModelTrainingAnnotationSpecification = {
+export type AiModelTrainingAnnotationSpecification = {
   annotationId: string;
   colorLayerName: string;
   segmentationLayerName: string;

frontend/javascripts/admin/voxelytics/ai_model_list_view.tsx

@@ -1,9 +1,9 @@
 import _ from "lodash";
 import React, { useState } from "react";
-import { SyncOutlined } from "@ant-design/icons";
-import { Table, Button, Modal } from "antd";
-import { getAiModels } from "admin/admin_rest_api";
-import type { AiModel } from "types/api_flow_types";
+import { PlusOutlined, SyncOutlined } from "@ant-design/icons";
+import { Table, Button, Modal, Space } from "antd";
+import { getAiModels, getTracingForAnnotationType } from "admin/admin_rest_api";
+import type { AiModel, APIAnnotation, ServerVolumeTracing } from "types/api_flow_types";
 import FormattedDate from "components/formatted_date";
 import { formatUserName } from "oxalis/model/accessors/user_accessor";
 import { useSelector } from "react-redux";
@@ -12,11 +12,19 @@ import { JobState } from "admin/job/job_list_view";
 import { Link } from "react-router-dom";
 import { useGuardedFetch } from "libs/react_helpers";
 import { PageNotAvailableToNormalUser } from "components/permission_enforcer";
+import { type AnnotationWithDataset, TrainAiModelTab } from "oxalis/view/jobs/train_ai_model";
+import {
+  getResolutionInfo,
+  getSegmentationLayerByName,
+} from "oxalis/model/accessors/dataset_accessor";
+import { serverVolumeToClientVolumeTracing } from "oxalis/model/reducers/volumetracing_reducer";
+import type { Vector3 } from "oxalis/constants";
 import type { Key } from "react";
 
 export default function AiModelListView() {
   const activeUser = useSelector((state: OxalisState) => state.activeUser);
   const [refreshCounter, setRefreshCounter] = useState(0);
+  const [isTrainModalVisible, setIsTrainModalVisible] = useState(false);
   const [aiModels, isLoading] = useGuardedFetch(
     getAiModels,
     [],
@@ -30,10 +38,18 @@ export default function AiModelListView() {
 
   return (
     <div className="container voxelytics-view">
+      {isTrainModalVisible ? (
+        <TrainNewAiJobModal onClose={() => setIsTrainModalVisible(false)} />
+      ) : null}
       <div className="pull-right">
-        <Button onClick={() => setRefreshCounter((val) => val + 1)}>
-          <SyncOutlined spin={isLoading} /> Refresh
-        </Button>
+        <Space>
+          <Button onClick={() => setIsTrainModalVisible(true)}>
+            <PlusOutlined /> Train new Model
+          </Button>
+          <Button onClick={() => setRefreshCounter((val) => val + 1)}>
+            <SyncOutlined spin={isLoading} /> Refresh
+          </Button>
+        </Space>
       </div>
       <h3>AI Models</h3>
       <Table
@@ -92,6 +108,69 @@ export default function AiModelListView() {
   );
 }
 
+function TrainNewAiJobModal({ onClose }: { onClose: () => void }) {
+  const [annotationsWithDatasets, setAnnotationsWithDatasets] = useState<
+    AnnotationWithDataset<APIAnnotation>[]
+  >([]);
+
+  const getMagForSegmentationLayer = async (annotationId: string, layerName: string) => {
+    // The layer name is a human-readable one. It can either belong to an annotationLayer
+    // (threfore, also to a volume tracing) or to the actual dataset.
+    // Both are checked below. This won't be ambiguous because annotationLayers must not
+    // have names that dataset layers already have.
+
+    const annotationWithDataset = annotationsWithDatasets.find(({ annotation }) => {
+      const currentAnnotationId = annotation.id;
+      return annotationId === currentAnnotationId;
+    });
+    if (annotationWithDataset == null) {
+      throw new Error("Cannot find annotation for specified id.");
+    }
+
+    const { annotation, dataset } = annotationWithDataset;
+
+    let annotationLayer = annotation.annotationLayers.find((l) => l.name === layerName);
+
+    if (annotationLayer != null) {
+      const volumeTracing = (await getTracingForAnnotationType(
+        annotation,
+        annotationLayer,
+      )) as ServerVolumeTracing;
+      const resolutions: Vector3[] = volumeTracing.resolutions?.map(({ x, y, z }) => [x, y, z]) || [
+        [1, 1, 1],
+      ];
+      return getResolutionInfo(resolutions).getFinestResolution();
+    } else {
+      const segmentationLayer = getSegmentationLayerByName(dataset, layerName);
+      return getResolutionInfo(segmentationLayer.resolutions).getFinestResolution();
+    }
+  };
+
+  return (
+    <Modal
+      width={875}
+      open
+      title={
+        <>
+          <i className="fas fa-magic icon-margin-right" />
+          AI Analysis
+        </>
+      }
+      onCancel={onClose}
+      footer={null}
+    >
+      <TrainAiModelTab
+        getMagForSegmentationLayer={getMagForSegmentationLayer}
+        onClose={onClose}
+        annotationsWithDatasets={annotationsWithDatasets}
+        onAddAnnotationsWithDatasets={(newItems) => {
+          setAnnotationsWithDatasets([...annotationsWithDatasets, ...newItems]);
+        }}
+      />
+    </Modal>
+  );
+}
+
 const renderActionsForModel = (model: AiModel) => {
   if (model.trainingJob == null) {
     return;

frontend/javascripts/components/layer_selection.tsx

@@ -2,27 +2,28 @@ import { Form, Select } from "antd";
 import { getReadableNameOfVolumeLayer } from "oxalis/model/accessors/volumetracing_accessor";
 import type { HybridTracing } from "oxalis/store";
 import type React from "react";
-import type { APIDataLayer } from "types/api_flow_types";
+import type { APIAnnotation, APIDataLayer } from "types/api_flow_types";
 
-type LayerSelectionProps = {
+type LayerSelectionProps<L extends { name: string }> = {
+  name: string | Array<string | number>;
   chooseSegmentationLayer: boolean;
-  layers: APIDataLayer[];
-  tracing: HybridTracing;
+  layers: L[];
+  getReadableNameForLayer: (layer: L) => string;
   fixedLayerName?: string;
   label?: string;
 };
 
-export function LayerSelection({
+export function LayerSelection<L extends { name: string }>({
   layers,
-  tracing,
+  getReadableNameForLayer,
   fixedLayerName,
   layerType,
   onChange,
   style,
   value,
 }: {
-  layers: APIDataLayer[];
-  tracing: HybridTracing;
+  layers: L[];
+  getReadableNameForLayer: (layer: L) => string;
   fixedLayerName?: string;
   layerType?: string;
   style?: React.CSSProperties;
@@ -49,7 +50,7 @@ export function LayerSelection({
       value={value}
     >
       {layers.map((layer) => {
-        const readableName = getReadableNameOfVolumeLayer(layer, tracing) || layer.name;
+        const readableName = getReadableNameForLayer(layer);
         return (
           <Select.Option key={layer.name} value={layer.name}>
             {readableName}
@@ -60,18 +61,19 @@ export function LayerSelection({
   );
 }
 
-export function LayerSelectionFormItem({
+export function LayerSelectionFormItem<L extends { name: string }>({
+  name,
   chooseSegmentationLayer,
   layers,
-  tracing,
+  getReadableNameForLayer,
   fixedLayerName,
   label,
-}: LayerSelectionProps): JSX.Element {
+}: LayerSelectionProps<L>): JSX.Element {
   const layerType = chooseSegmentationLayer ? "segmentation" : "color";
   return (
     <Form.Item
       label={label || "Layer"}
-      name="layerName"
+      name={name}
       rules={[
         {
           required: true,
@@ -85,7 +87,7 @@ export function LayerSelectionFormItem({
         layers={layers}
         fixedLayerName={fixedLayerName}
         layerType={layerType}
-        tracing={tracing}
+        getReadableNameForLayer={getReadableNameForLayer}
       />
     </Form.Item>
   );

frontend/javascripts/oxalis/default_state.ts

@@ -240,7 +240,7 @@ const defaultState: OxalisState = {
     showVersionRestore: false,
     showDownloadModal: false,
     showPythonClientModal: false,
-    aIJobModalState: "invisible",
+    aIJobModalState: "neuron_inferral",
     showRenderAnimationModal: false,
     showShareModal: false,
     storedLayouts: {},

frontend/javascripts/oxalis/model/accessors/volumetracing_accessor.ts

@@ -118,7 +118,32 @@ export function getReadableNameByVolumeTracingId(
   tracingId: string,
 ) {
   const volumeDescriptor = getVolumeDescriptorById(annotation, tracingId);
-  return volumeDescriptor.name || "Volume";
+  return volumeDescriptor.name;
+}
+
+export function getSegmentationLayerByHumanReadableName(
+  dataset: APIDataset,
+  annotation: APIAnnotation | HybridTracing,
+  name: string,
+) {
+  try {
+    const layer = getSegmentationLayerByName(dataset, name);
+    return layer;
+  } catch {}
+
+  const layer = getVolumeTracingLayers(dataset).find((currentLayer) => {
+    if (currentLayer.tracingId == null) {
+      throw new Error("getVolumeTracingLayers must return tracing.");
+    }
+    const readableName = getReadableNameByVolumeTracingId(annotation, currentLayer.tracingId);
+    return readableName === name;
+  });
+
+  if (layer == null) {
+    throw new Error("Could not find segmentation layer with the name: " + name);
+  }
+
+  return layer;
 }
 
 export function getAllReadableLayerNames(dataset: APIDataset, tracing: Tracing) {
@@ -135,7 +160,7 @@ export function getAllReadableLayerNames(dataset: APIDataset, tracing: Tracing)
 
 export function getReadableNameForLayerName(
   dataset: APIDataset,
-  tracing: Tracing,
+  tracing: APIAnnotation | HybridTracing,
   layerName: string,
 ): string {
   const layer = getLayerByName(dataset, layerName, true);
@@ -842,7 +867,7 @@ export const getBucketRetrievalSourceFn =
 
 export function getReadableNameOfVolumeLayer(
   layer: APIDataLayer,
-  tracing: HybridTracing,
+  tracing: APIAnnotation | HybridTracing,
 ): string | null {
   return "tracingId" in layer && layer.tracingId != null
     ? getReadableNameByVolumeTracingId(tracing, layer.tracingId)

frontend/javascripts/oxalis/model_initialization.ts

@@ -401,7 +401,7 @@ function initializeDataset(
   const volumeTracings = getServerVolumeTracings(serverTracings);
 
   if (volumeTracings.length > 0) {
-    const newDataLayers = setupLayerForVolumeTracing(dataset, volumeTracings);
+    const newDataLayers = getMergedDataLayersFromDatasetAndVolumeTracings(dataset, volumeTracings);
     mutableDataset.dataSource.dataLayers = newDataLayers;
     validateVolumeLayers(volumeTracings, newDataLayers);
   }
@@ -480,7 +480,7 @@ function initializeDataLayerInstances(gpuFactor: number | null | undefined): {
   };
 }
 
-function setupLayerForVolumeTracing(
+function getMergedDataLayersFromDatasetAndVolumeTracings(
   dataset: APIDataset,
   tracings: Array<ServerVolumeTracing>,
 ): Array<APIDataLayer> {

frontend/javascripts/oxalis/view/action-bar/create_animation_modal.tsx

@@ -94,7 +94,6 @@ export default function CreateAnimationModalWrapper(props: Props) {
 function CreateAnimationModal(props: Props) {
   const { isOpen, onClose } = props;
   const dataset = useSelector((state: OxalisState) => state.dataset);
-  const tracing = useSelector((state: OxalisState) => state.tracing);
   const activeOrganization = useSelector((state: OxalisState) => state.activeOrganization);
 
   const colorLayers = getColorLayers(dataset);
@@ -403,7 +402,7 @@ function CreateAnimationModal(props: Props) {
               layers={colorLayers}
               value={selectedColorLayerName}
               onChange={setSelectedColorLayerName}
-              tracing={tracing}
+              getReadableNameForLayer={(layer) => layer.name}
               style={{ width: "100%" }}
             />
           </Col>

frontend/javascripts/oxalis/view/action-bar/download_modal_view.tsx

@@ -630,7 +630,9 @@ function _DownloadModalView({
             layers={layers}
             value={selectedLayerName}
             onChange={setSelectedLayerName}
-            tracing={tracing}
+            getReadableNameForLayer={(layer) =>
+              getReadableNameOfVolumeLayer(layer, tracing) || layer.name
+            }
             style={{ width: "100%" }}
           />
 

frontend/javascripts/oxalis/view/action-bar/starting_job_modals.tsx

@@ -51,7 +51,10 @@ import { isBoundingBoxExportable } from "./download_modal_view";
 import features from "features";
 import { setAIJobModalStateAction } from "oxalis/model/actions/ui_actions";
 import { InfoCircleOutlined } from "@ant-design/icons";
-import { TrainAiModelTab, CollapsibleWorkflowYamlEditor } from "../jobs/train_ai_model";
+import {
+  CollapsibleWorkflowYamlEditor,
+  TrainAiModelFromAnnotationTab,
+} from "../jobs/train_ai_model";
 import { LayerSelectionFormItem } from "components/layer_selection";
 import { useGuardedFetch } from "libs/react_helpers";
 import _ from "lodash";
@@ -291,7 +294,7 @@ export function StartAIJobModal({ aIJobModalState }: StartAIJobModalProps) {
       ? {
           label: "Train a model",
           key: "trainModel",
-          children: <TrainAiModelTab onClose={onClose} />,
+          children: <TrainAiModelFromAnnotationTab onClose={onClose} />,
         }
       : null,
     isSuperUser
@@ -627,11 +630,14 @@ function StartJobForm(props: StartJobFormProps) {
         initialName={`${dataset.name}_${props.suggestedDatasetSuffix}`}
       />
       <LayerSelectionFormItem
+        name="layerName"
         chooseSegmentationLayer={chooseSegmentationLayer}
         label={chooseSegmentationLayer ? "Segmentation Layer" : "Image data layer"}
         layers={layers}
         fixedLayerName={fixedSelectedLayer?.name}
-        tracing={tracing}
+        getReadableNameForLayer={(layer) =>
+          getReadableNameOfVolumeLayer(layer, tracing) || layer.name
+        }
       />
       <BoundingBoxSelectionFormItem
         isBoundingBoxConfigurable={isBoundingBoxConfigurable}