Add nd support for tiff export job (#7971)

* Add nd support for tiff export job * remove nd warning from tiff export modal * add nd info to zarr2 streaming metadata * add changelog entry * transform input bbox & additionalCoordinates to nd bounding box for worker input * format backend * apply feedback * fix rendering of new nd bounding box tiff export job arg in jobs list --------- Co-authored-by: Michael Büßemeyer <[email protected]> Co-authored-by: Florian M <[email protected]> Co-authored-by: Florian M <[email protected]>
scalableminds · Sep 19, 2024 · 3d2a88e · 3d2a88e
1 parent d336b7c
commit 3d2a88e
Show file tree

Hide file tree

Showing 14 changed files with 130 additions and 30 deletions.
diff --git a/CHANGELOG.unreleased.md b/CHANGELOG.unreleased.md
@@ -12,6 +12,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 
 ### Added
 - It is now possible to focus a bounding box in the bounding box tab by clicking its edges in a viewport or via a newly added context menu entry. [#8054](https://github.com/scalableminds/webknossos/pull/8054)
+- Added the option to export nd datasets as ome tiff or tiff stack. Previously, this was only possible for 3d datasets. [#7971](https://github.com/scalableminds/webknossos/pull/7971)
 - Added an assertion to the backend to ensure unique keys in the metadata info of datasets and folders. [#8068](https://github.com/scalableminds/webknossos/issues/8068)
 
 ### Changed

diff --git a/app/controllers/JobController.scala b/app/controllers/JobController.scala
@@ -4,7 +4,7 @@ import play.silhouette.api.Silhouette
 import com.scalableminds.util.geometry.{BoundingBox, Vec3Double, Vec3Int}
 import com.scalableminds.util.accesscontext.GlobalAccessContext
 import com.scalableminds.util.tools.Fox
-import models.dataset.{DataStoreDAO, DatasetDAO, DatasetService}
+import models.dataset.{DataStoreDAO, DatasetDAO, DatasetLayerAdditionalAxesDAO, DatasetService}
 import models.job._
 import models.organization.OrganizationDAO
 import models.user.MultiUserDAO
@@ -19,7 +19,9 @@ import java.util.Date
 import javax.inject.Inject
 import scala.concurrent.ExecutionContext
 import com.scalableminds.util.enumeration.ExtendedEnumeration
-import com.scalableminds.webknossos.datastore.models.{LengthUnit, VoxelSize}
+import com.scalableminds.webknossos.datastore.dataformats.zarr.Zarr3OutputHelper
+import com.scalableminds.webknossos.datastore.datareaders.{AxisOrder, FullAxisOrder, NDBoundingBox}
+import com.scalableminds.webknossos.datastore.models.{AdditionalCoordinate, LengthUnit, VoxelSize}
 import models.team.PricingPlan
 
 object MovieResolutionSetting extends ExtendedEnumeration {
@@ -54,13 +56,15 @@ class JobController @Inject()(
     jobService: JobService,
     workerService: WorkerService,
     workerDAO: WorkerDAO,
+    datasetLayerAdditionalAxesDAO: DatasetLayerAdditionalAxesDAO,
     wkconf: WkConf,
     multiUserDAO: MultiUserDAO,
     wkSilhouetteEnvironment: WkSilhouetteEnvironment,
     slackNotificationService: SlackNotificationService,
     organizationDAO: OrganizationDAO,
     dataStoreDAO: DataStoreDAO)(implicit ec: ExecutionContext, playBodyParsers: PlayBodyParsers)
-    extends Controller {
+    extends Controller
+    with Zarr3OutputHelper {
 
   def status: Action[AnyContent] = sil.SecuredAction.async { implicit request =>
     for {
@@ -316,6 +320,7 @@ class JobController @Inject()(
   def runExportTiffJob(organizationId: String,
                        datasetName: String,
                        bbox: String,
+                       additionalCoordinates: Option[String],
                        layerName: Option[String],
                        mag: Option[String],
                        annotationLayerName: Option[String],
@@ -330,6 +335,20 @@ class JobController @Inject()(
           _ <- Fox.runOptional(layerName)(datasetService.assertValidLayerNameLax)
           _ <- Fox.runOptional(annotationLayerName)(datasetService.assertValidLayerNameLax)
           _ <- jobService.assertBoundingBoxLimits(bbox, mag)
+          additionalAxesOpt <- Fox.runOptional(layerName)(layerName =>
+            datasetLayerAdditionalAxesDAO.findAllForDatasetAndDataLayerName(dataset._id, layerName))
+          additionalAxesOpt <- Fox.runOptional(additionalAxesOpt)(a => Fox.successful(reorderAdditionalAxes(a)))
+          rank = additionalAxesOpt.map(_.length).getOrElse(0) + 4
+          axisOrder = FullAxisOrder.fromAxisOrderAndAdditionalAxes(rank,
+                                                                   AxisOrder.cAdditionalxyz(rank),
+                                                                   additionalAxesOpt)
+          threeDBBox <- BoundingBox.fromLiteral(bbox).toFox ~> "job.invalidBoundingBox"
+          parsedAdditionalCoordinatesOpt <- Fox.runOptional(additionalCoordinates)(coords =>
+            Json.parse(coords).validate[Seq[AdditionalCoordinate]]) ~> "job.additionalCoordinates.invalid"
+          parsedAdditionalCoordinates = parsedAdditionalCoordinatesOpt.getOrElse(Seq.empty)
+          additionalAxesOfNdBBox = additionalAxesOpt.map(additionalAxes =>
+            additionalAxes.map(_.intersectWithAdditionalCoordinates(parsedAdditionalCoordinates)))
+          ndBoundingBox = NDBoundingBox(threeDBBox, additionalAxesOfNdBBox.getOrElse(Seq.empty), axisOrder)
           command = JobCommand.export_tiff
           exportFileName = if (asOmeTiff)
             s"${formatDateForFilename(new Date())}__${datasetName}__${annotationLayerName.map(_ => "volume").getOrElse(layerName.getOrElse(""))}.ome.tif"
@@ -338,12 +357,12 @@ class JobController @Inject()(
           commandArgs = Json.obj(
             "organization_name" -> organizationId,
             "dataset_name" -> datasetName,
-            "bbox" -> bbox,
+            "nd_bbox" -> ndBoundingBox.toWkLibsDict,
             "export_file_name" -> exportFileName,
             "layer_name" -> layerName,
             "mag" -> mag,
             "annotation_layer_name" -> annotationLayerName,
-            "annotation_id" -> annotationId
+            "annotation_id" -> annotationId,
           )
           job <- jobService.submitJob(command, commandArgs, request.identity, dataset._dataStore) ?~> "job.couldNotRunTiffExport"
           js <- jobService.publicWrites(job)

diff --git a/conf/messages b/conf/messages
@@ -335,6 +335,7 @@ job.emailNotifactionsDisabled = Email notifications are not enabled for this job
 job.renderAnimation.notAllowed.organization = "Rendering animations is only allowed for datasets of your own organization."
 job.alignSections.notAllowed.organization = "Aligning sections is only allowed for datasets of your own organization."
 job.alignSections.notAllowed.onlySuperUsers = "For now, aligning sections is only allowed for super users."
+job.additionalCoordinates.invalid = "The passed additional coordinates are invalid."
 
 voxelytics.disabled = Voxelytics workflow reporting and logging are not enabled for this WEBKNOSSOS instance.
 voxelytics.runNotFound = Workflow runs not found

diff --git a/conf/webknossos.latest.routes b/conf/webknossos.latest.routes
@@ -262,7 +262,7 @@ GET           /jobs/status
 POST          /jobs/run/convertToWkw/:organizationId/:datasetName                               controllers.JobController.runConvertToWkwJob(organizationId: String, datasetName: String, scale: String, unit: Option[String])
 POST          /jobs/run/computeMeshFile/:organizationId/:datasetName                            controllers.JobController.runComputeMeshFileJob(organizationId: String, datasetName: String, layerName: String, mag: String, agglomerateView: Option[String])
 POST          /jobs/run/computeSegmentIndexFile/:organizationId/:datasetName                    controllers.JobController.runComputeSegmentIndexFileJob(organizationId: String, datasetName: String, layerName: String)
-POST          /jobs/run/exportTiff/:organizationId/:datasetName                                 controllers.JobController.runExportTiffJob(organizationId: String, datasetName: String, bbox: String, layerName: Option[String], mag: Option[String], annotationLayerName: Option[String], annotationId: Option[String], asOmeTiff: Boolean)
+POST          /jobs/run/exportTiff/:organizationId/:datasetName                                 controllers.JobController.runExportTiffJob(organizationId: String, datasetName: String, bbox: String, additionalCoordinates: Option[String], layerName: Option[String], mag: Option[String], annotationLayerName: Option[String], annotationId: Option[String], asOmeTiff: Boolean)
 POST          /jobs/run/inferNuclei/:organizationId/:datasetName                                controllers.JobController.runInferNucleiJob(organizationId: String, datasetName: String, layerName: String, newDatasetName: String)
 POST          /jobs/run/inferNeurons/:organizationId/:datasetName                               controllers.JobController.runInferNeuronsJob(organizationId: String, datasetName: String, layerName: String, bbox: String, newDatasetName: String)
 POST          /jobs/run/inferMitochondria/:organizationId/:datasetName                          controllers.JobController.runInferMitochondriaJob(organizationId: String, datasetName: String, layerName: String, bbox: String, newDatasetName: String)

diff --git a/frontend/javascripts/admin/api/jobs.ts b/frontend/javascripts/admin/api/jobs.ts
@@ -9,6 +9,7 @@ import type {
   APIEffectiveJobState,
   AiModel,
   RenderAnimationOptions,
+  AdditionalCoordinate,
 } from "types/api_flow_types";
 import { assertResponseLimit } from "./api_utils";
 
@@ -21,6 +22,7 @@ function transformBackendJobToAPIJob(job: any): APIJob {
     layerName: job.commandArgs.layer_name || job.commandArgs.volume_layer_name,
     annotationLayerName: job.commandArgs.annotation_layer_name,
     boundingBox: job.commandArgs.bbox,
+    ndBoundingBox: job.commandArgs.nd_bbox,
     exportFileName: job.commandArgs.export_file_name,
     tracingId: job.commandArgs.volume_tracing_id,
     annotationId: job.commandArgs.annotation_id,
@@ -100,6 +102,7 @@ export async function startExportTiffJob(
   datasetName: string,
   organizationId: string,
   bbox: Vector6,
+  additionalCoordinates: AdditionalCoordinate[] | null,
   layerName: string | null | undefined,
   mag: string | null | undefined,
   annotationId: string | null | undefined,
@@ -119,6 +122,9 @@ export async function startExportTiffJob(
   if (annotationLayerName != null) {
     params.append("annotationLayerName", annotationLayerName);
   }
+  if (additionalCoordinates != null) {
+    params.append("additionalCoordinates", JSON.stringify(additionalCoordinates));
+  }
   return Request.receiveJSON(
     `/api/jobs/run/exportTiff/${organizationId}/${datasetName}?${params}`,
     {

diff --git a/frontend/javascripts/admin/job/job_list_view.tsx b/frontend/javascripts/admin/job/job_list_view.tsx
@@ -23,6 +23,7 @@ import FormattedDate from "components/formatted_date";
 import { AsyncLink } from "components/async_clickables";
 import { useEffect, useState } from "react";
 import { useInterval } from "libs/react_helpers";
+import { formatWkLibsNdBBox } from "libs/format_utils";
 
 // Unfortunately, the twoToneColor (nor the style) prop don't support
 // CSS variables.
@@ -129,7 +130,7 @@ function JobListView() {
       return (
         <span>
           Tiff export of layer {layerLabel} from {labelToAnnotationOrDataset} (Bounding Box{" "}
-          {job.boundingBox})
+          {job.ndBoundingBox ? formatWkLibsNdBBox(job.ndBoundingBox) : job.boundingBox})
         </span>
       );
     } else if (job.type === APIJobType.RENDER_ANIMATION && job.organizationId && job.datasetName) {

diff --git a/frontend/javascripts/libs/format_utils.ts b/frontend/javascripts/libs/format_utils.ts
@@ -16,7 +16,7 @@ import localeData from "dayjs/plugin/localeData";
 
 import type { BoundingBoxObject } from "oxalis/store";
 import type { Duration } from "dayjs/plugin/duration";
-import type { VoxelSize } from "types/api_flow_types";
+import type { VoxelSize, WkLibsNdBoundingBox } from "types/api_flow_types";
 
 dayjs.extend(updateLocale);
 dayjs.extend(duration);
@@ -537,3 +537,18 @@ export function formatVoxels(voxelCount: number) {
 export function formatNumber(num: number): string {
   return new Intl.NumberFormat("en-US").format(num);
 }
+
+export function formatWkLibsNdBBox(ndBBox: WkLibsNdBoundingBox): string {
+  let bboxString = ndBBox.topLeft.join(", ");
+  bboxString += ` ${ndBBox.width}, ${ndBBox.height}, ${ndBBox.depth}`;
+  let additionalAxisStrings = ndBBox.additionalAxes.map((axis) => {
+    const boundsString =
+      axis.bounds[1] - 1 > axis.bounds[0]
+        ? `${axis.bounds[0]}-${axis.bounds[1]}`
+        : `${axis.bounds[0]}`;
+    return `${axis.name}=${boundsString}`;
+  });
+  return additionalAxisStrings.length > 0
+    ? `${bboxString} (${additionalAxisStrings.join(", ")})`
+    : bboxString;
+}
diff --git a/frontend/javascripts/oxalis/view/action-bar/download_modal_view.tsx b/frontend/javascripts/oxalis/view/action-bar/download_modal_view.tsx
@@ -58,6 +58,7 @@ import { formatCountToDataAmountUnit, formatScale } from "libs/format_utils";
 import type { BoundingBoxType, Vector3 } from "oxalis/constants";
 import { useStartAndPollJob } from "admin/job/job_hooks";
 import { LayerSelection } from "components/layer_selection";
+import { getAdditionalCoordinatesAsString } from "oxalis/model/accessors/flycam_accessor";
 const { Paragraph, Text } = Typography;
 
 type TabKeys = "download" | "export" | "python";
@@ -281,6 +282,9 @@ function _DownloadModalView({
   const rawUserBoundingBoxes = useSelector((state: OxalisState) =>
     getUserBoundingBoxesFromState(state),
   );
+  const currentAdditionalCoordinates = useSelector(
+    (state: OxalisState) => state.flycam.additionalCoordinates,
+  );
   const typeName = isAnnotation ? "annotation" : "dataset";
   const isMergerModeEnabled = useSelector(
     (state: OxalisState) => state.temporaryConfiguration.isMergerModeEnabled,
@@ -357,16 +361,13 @@ function _DownloadModalView({
       );
       onClose();
     } else if (activeTabKey === "export" && startJob != null) {
-      if ((selectedLayerInfos.additionalAxes || []).length > 0) {
-        Toast.warning("Exporting an n-dimensional layer is currently not supported.");
-        return;
-      }
       await Model.ensureSavedState();
       await startJob(async () => {
         const job = await startExportTiffJob(
           dataset.name,
           dataset.owningOrganization,
           computeArrayFromBoundingBox(selectedBoundingBox.boundingBox),
+          currentAdditionalCoordinates,
           selectedLayerInfos.layerName,
           mag.join("-"),
           selectedLayerInfos.annotationId,
@@ -396,18 +397,6 @@ function _DownloadModalView({
           </Text>
         </Row>
       ) : null;
-    const ndVolumeWarning = isVolumeNDimensional ? (
-      <Row key="unsupported-nd">
-        <Text
-          style={{
-            margin: "0 6px 12px",
-          }}
-          type="warning"
-        >
-          Downloading/exporting n-dimensional volume data is not yet supported.
-        </Text>
-      </Row>
-    ) : null;
     const pythonTokenWarning =
       activeTabKey === "python" ? (
         <Row key="python-token-warning">
@@ -424,7 +413,7 @@ function _DownloadModalView({
         </Row>
       ) : null;
 
-    return [volumeFallbackWarning, ndVolumeWarning, pythonTokenWarning];
+    return [volumeFallbackWarning, pythonTokenWarning];
   };
 
   const handleTabChange = (key: string) => {
@@ -652,6 +641,26 @@ function _DownloadModalView({
             style={{ width: "100%" }}
           />
           {boundingBoxCompatibilityAlerts}
+          {selectedLayerInfos.additionalAxes != null && (
+            <Row>
+              <Divider
+                style={{
+                  margin: "18px 0",
+                }}
+              >
+                Additional Coordinates
+              </Divider>
+              <Text
+                style={{
+                  margin: "0 6px 12px",
+                }}
+              >
+                Your dataset has more than three dimensions. The export will only include the
+                selected bounding box at the current additional dimensions:{" "}
+                {getAdditionalCoordinatesAsString(currentAdditionalCoordinates)}
+              </Text>
+            </Row>
+          )}
 
           <Divider
             style={{

diff --git a/frontend/javascripts/types/api_flow_types.ts b/frontend/javascripts/types/api_flow_types.ts
@@ -694,6 +694,11 @@ export enum APIJobType {
   INFER_MITOCHONDRIA = "infer_mitochondria",
 }
 
+export type WkLibsNdBoundingBox = BoundingBoxObject & {
+  axisOrder: { string: number };
+  additionalAxes: Array<AdditionalAxis>;
+};
+
 export type APIJob = {
   readonly id: string;
   readonly datasetName: string | null | undefined;
@@ -705,6 +710,7 @@ export type APIJob = {
   readonly annotationType: string | null | undefined;
   readonly organizationId: string | null | undefined;
   readonly boundingBox: string | null | undefined;
+  readonly ndBoundingBox: WkLibsNdBoundingBox | null | undefined;
   readonly mergeSegments: boolean | null | undefined;
   readonly type: APIJobType;
   readonly state: APIEffectiveJobState;

diff --git a/util/src/main/scala/com/scalableminds/util/geometry/BoundingBox.scala b/util/src/main/scala/com/scalableminds/util/geometry/BoundingBox.scala
@@ -2,6 +2,7 @@ package com.scalableminds.util.geometry
 
 import com.scalableminds.util.tools.Math.ceilDiv
 import net.liftweb.common.Full
+import play.api.libs.json.{JsObject, Json}
 
 case class BoundingBox(topLeft: Vec3Int, width: Int, height: Int, depth: Int) {
 
@@ -70,6 +71,9 @@ case class BoundingBox(topLeft: Vec3Int, width: Int, height: Int, depth: Int) {
     Vec3Int(width, height, depth)
 
   def toLiteral: String = f"${topLeft.x},${topLeft.y},${topLeft.z},$width,$height,$depth"
+
+  def toWkLibsJson: JsObject =
+    Json.obj("topLeft" -> topLeft, "width" -> width, "height" -> height, "depth" -> depth)
 }
 
 object BoundingBox {

diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/AxisOrder.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/AxisOrder.scala
@@ -1,7 +1,7 @@
 package com.scalableminds.webknossos.datastore.datareaders
 
 import com.scalableminds.webknossos.datastore.models.datasource.AdditionalAxis
-import play.api.libs.json.{Json, OFormat}
+import play.api.libs.json.{JsValue, Json, OFormat}
 
 // Defines the axis order of a DatasetArray. Note that this ignores transpose codecs/ArrayOrder.F/C.
 // Those will have to be applied on individual chunk’s contents.
@@ -97,6 +97,12 @@ case class FullAxisOrder(axes: Seq[Axis]) {
   def permuteIndicesArrayToWk(indices: Array[Int]): Array[Int] =
     arrayToWkPermutation.map(indices(_))
 
+  def toWkLibsJson: JsValue =
+    Json.toJson(axes.zipWithIndex.collect {
+      case (axis, index) if axis.name == "x" || axis.name == "y" || axis.name == "z" =>
+        axis.name -> index
+    }.toMap)
+
 }
 
 object FullAxisOrder {

diff --git a/...ssos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/NDBoundingBox.scala b/...ssos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/NDBoundingBox.scala
@@ -0,0 +1,15 @@
+package com.scalableminds.webknossos.datastore.datareaders
+
+import com.scalableminds.util.geometry.BoundingBox
+import com.scalableminds.webknossos.datastore.models.datasource.AdditionalAxis
+import play.api.libs.json.{JsValue, Json}
+
+case class NDBoundingBox(boundingBox: BoundingBox, additionalAxes: Seq[AdditionalAxis], fullAxisOrder: FullAxisOrder) {
+
+  def toWkLibsDict: JsValue = {
+    val additionalAxesDict = Json.toJson(additionalAxes)
+    val axisOrderDict = fullAxisOrder.toWkLibsJson
+    boundingBox.toWkLibsJson ++ Json.obj("additionalAxes" -> additionalAxesDict, "axisOrder" -> axisOrderDict)
+  }
+
+}
diff --git a/...os-datastore/app/com/scalableminds/webknossos/datastore/datareaders/zarr/ZarrHeader.scala b/...os-datastore/app/com/scalableminds/webknossos/datastore/datareaders/zarr/ZarrHeader.scala
@@ -62,15 +62,19 @@ object ZarrHeader extends JsonImplicits {
     // data request method always decompresses before sending
     val compressor = None
 
-    val shape = Array(
-      channels,
+    val additionalAxesShapeEntries =
+      dataLayer.additionalAxes.map(axes => axes.map(_.bounds(1)).toArray).getOrElse(Array.empty)
+    val additionalAxesChunksEntries =
+      dataLayer.additionalAxes.map(axes => axes.map(_ => 1).toArray).getOrElse(Array.empty)
+
+    val shape = Array(channels) ++ additionalAxesShapeEntries ++ Array(
       // Zarr can't handle data sets that don't start at 0, so we extend the shape to include "true" coords
       (dataLayer.boundingBox.width + dataLayer.boundingBox.topLeft.x) / mag.x,
       (dataLayer.boundingBox.height + dataLayer.boundingBox.topLeft.y) / mag.y,
       (dataLayer.boundingBox.depth + dataLayer.boundingBox.topLeft.z) / mag.z
     )
 
-    val chunks = Array(channels, cubeLength, cubeLength, cubeLength)
+    val chunks = Array(channels) ++ additionalAxesChunksEntries ++ Array(cubeLength, cubeLength, cubeLength)
 
     ZarrHeader(zarr_format = 2,
                shape = shape,