Remove VoxelType Enum, use ArrayDataType instead for WKW (#8559)

fm3 · web-flow · commit 2da75a5271e9 · 2025-05-12T10:16:41.000+02:00
* WIP: unify backend dtype classes

* Fix VoxelType off-by-one

* fix merge artifact in migrations guide
diff --git a/MIGRATIONS.unreleased.md b/MIGRATIONS.unreleased.md
@@ -9,7 +9,6 @@ User-facing changes are documented in the [changelog](CHANGELOG.released.md).
 [Commits](https://github.com/scalableminds/webknossos/compare/25.05.0...HEAD)
 
 ### Postgres Evolutions:
-<<<<<<< HEAD
 
 - [131-more-indices-on-users.sql](conf/evolutions/131-more-indices-on-users.sql)
 - [132-remove-stored-meshes.sql](conf/evolutions/132-remove-stored-meshes.sql)
diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/dataformats/wkw/VoxelType.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/dataformats/wkw/VoxelType.scala
diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/dataformats/wkw/WKWHeader.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/dataformats/wkw/WKWHeader.scala
@@ -7,6 +7,7 @@ import com.scalableminds.webknossos.datastore.datareaders.ArrayDataType.ArrayDat
 import com.scalableminds.webknossos.datastore.datareaders.ArrayOrder.ArrayOrder
 import com.scalableminds.webknossos.datastore.datareaders.DimensionSeparator.DimensionSeparator
 import com.scalableminds.webknossos.datastore.datareaders.{
+  ArrayDataType,
   ArrayOrder,
   Compressor,
   DatasetHeader,
@@ -32,8 +33,8 @@ case class WKWHeader(
     numChunksPerShardDimension: Int,
     numVoxelsPerChunkDimension: Int,
     blockType: ChunkType.Value,
-    voxelType: VoxelType.Value,
-    numBytesPerVoxel: Int, // this encodes the channel count together with voxelType
+    dataType: ArrayDataType,
+    numChannels: Int,
     jumpTable: Array[Long]
 ) extends DatasetHeader
     with WKWDataFormatHelper {
@@ -44,11 +45,11 @@ case class WKWHeader(
 
   def numChunksPerShard: Int = numChunksPerShardDimension * numChunksPerShardDimension * numChunksPerShardDimension
 
+  private def numBytesPerVoxel = numChannels * ArrayDataType.bytesPerElement(dataType)
+
   def numBytesPerChunk: Int =
     numVoxelsPerChunkDimension * numVoxelsPerChunkDimension * numVoxelsPerChunkDimension * numBytesPerVoxel
 
-  def numChannels: Int = numBytesPerVoxel / VoxelType.bytesPerVoxelPerChannel(voxelType)
-
   def blockLengths: Iterator[Int] =
     if (isCompressed) {
       jumpTable.sliding(2).map(a => (a(1) - a(0)).toInt)
@@ -64,7 +65,7 @@ case class WKWHeader(
     val sideLengths = (numChunksPerShardDimensionLog2 << 4) + numVoxelsPerChunkDimensionLog2
     output.writeByte(sideLengths)
     output.writeByte(blockType.id)
-    output.writeByte(voxelType.id)
+    output.writeByte(ArrayDataType.toWKWId(dataType))
     output.writeByte(numBytesPerVoxel)
     if (isHeaderFile) {
       output.writeLong(0L)
@@ -96,7 +97,7 @@ case class WKWHeader(
 
   override def order: ArrayOrder = ArrayOrder.F
 
-  override def resolvedDataType: ArrayDataType = VoxelType.toArrayDataType(voxelType)
+  override def resolvedDataType: ArrayDataType = dataType
 
   override def compressorImpl: Compressor = blockType match {
     case ChunkType.Raw                   => nullCompressor
@@ -146,7 +147,8 @@ object WKWHeader {
                                                                  numVoxelsPerChunkDimension,
                                                                  "[0, 1024)")
       blockType <- tryo(ChunkType(blockTypeId)) ?~! error("Specified blockType is not supported")
-      voxelType <- tryo(VoxelType(voxelTypeId)) ?~! error("Specified voxelType is not supported")
+      voxelType <- tryo(ArrayDataType.fromWKWTypeId(voxelTypeId)) ?~! error("Specified voxelType is not supported")
+      numChannels = numBytesPerVoxel / ArrayDataType.bytesPerElement(voxelType)
     } yield {
       val jumpTable = if (ChunkType.isCompressed(blockType) && readJumpTable) {
         val numChunksPerShard = numChunksPerShardDimension * numChunksPerShardDimension * numChunksPerShardDimension
@@ -159,7 +161,7 @@ object WKWHeader {
                     numVoxelsPerChunkDimension,
                     blockType,
                     voxelType,
-                    numBytesPerVoxel,
+                    numChannels,
                     jumpTable)
     }
   }
@@ -170,15 +172,15 @@ object WKWHeader {
   def apply(numChunksPerShardDimension: Int,
             numVoxelsPerChunkDimension: Int,
             blockType: ChunkType.Value,
-            voxelType: VoxelType.Value,
+            dataType: ArrayDataType,
             numChannels: Int): WKWHeader =
     new WKWHeader(
       currentVersion,
       numChunksPerShardDimension,
       numVoxelsPerChunkDimension,
       blockType,
-      voxelType,
-      VoxelType.bytesPerVoxelPerChannel(voxelType) * numChannels,
+      dataType,
+      numChannels,
       Array(0)
     )
 }
diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ArrayDataType.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ArrayDataType.scala
@@ -6,7 +6,7 @@ object ArrayDataType extends ExtendedEnumeration {
   type ArrayDataType = Value
   val f8, f4, i8, u8, i4, u4, i2, u2, i1, u1 = Value
 
-  def bytesPerElementFor(dataType: ArrayDataType): Int =
+  def bytesPerElement(dataType: ArrayDataType): Int =
     dataType match {
       case ArrayDataType.f8 => 8
       case ArrayDataType.f4 => 4
@@ -47,4 +47,32 @@ object ArrayDataType extends ExtendedEnumeration {
       case ArrayDataType.i1 => Byte.MinValue
       case ArrayDataType.u1 => 0
     }
+
+  def toWKWId(dataType: ArrayDataType): Int =
+    dataType match {
+      case ArrayDataType.u1 => 1
+      case ArrayDataType.u2 => 2
+      case ArrayDataType.u4 => 3
+      case ArrayDataType.u8 => 4
+      case ArrayDataType.f4 => 5
+      case ArrayDataType.f8 => 6
+      case ArrayDataType.i1 => 7
+      case ArrayDataType.i2 => 8
+      case ArrayDataType.i4 => 9
+      case ArrayDataType.i8 => 10
+    }
+
+  def fromWKWTypeId(wkwVoxelTypeId: Int): ArrayDataType.Value =
+    wkwVoxelTypeId match {
+      case 1  => ArrayDataType.u1
+      case 2  => ArrayDataType.u2
+      case 3  => ArrayDataType.u4
+      case 4  => ArrayDataType.u8
+      case 5  => ArrayDataType.f4
+      case 6  => ArrayDataType.f8
+      case 7  => ArrayDataType.i1
+      case 8  => ArrayDataType.i2
+      case 9  => ArrayDataType.i4
+      case 10 => ArrayDataType.i8
+    }
 }
diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/BytesConverter.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/BytesConverter.scala
@@ -1,6 +1,6 @@
 package com.scalableminds.webknossos.datastore.datareaders
 
-import com.scalableminds.webknossos.datastore.datareaders.ArrayDataType.{ArrayDataType, bytesPerElementFor}
+import com.scalableminds.webknossos.datastore.datareaders.ArrayDataType.ArrayDataType
 import net.liftweb.common.Box
 import net.liftweb.common.Box.tryo
 import ucar.ma2.{Array => MultiArray}
@@ -10,7 +10,7 @@ import java.nio.{ByteBuffer, ByteOrder}
 object BytesConverter {
   def toByteArray(multiArray: MultiArray, dataType: ArrayDataType, byteOrder: ByteOrder): Box[Array[Byte]] = tryo {
     val array = multiArray.getStorage
-    val bytesPerElement = bytesPerElementFor(dataType)
+    val bytesPerElement = ArrayDataType.bytesPerElement(dataType)
     // If the multiArray dtype size is 1, use the array directly.
     // This may be happen due to the skipTyping shortcut even for non-uint8-datasets
     if (multiArray.getDataType.getSize == 1) {
diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/DatasetHeader.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/DatasetHeader.scala
@@ -3,7 +3,7 @@ package com.scalableminds.webknossos.datastore.datareaders
 import com.scalableminds.util.geometry.{BoundingBox, Vec3Int}
 import com.scalableminds.webknossos.datastore.datareaders.ArrayOrder.ArrayOrder
 import com.scalableminds.webknossos.datastore.datareaders.DimensionSeparator.DimensionSeparator
-import ArrayDataType.{ArrayDataType, bytesPerElementFor}
+import ArrayDataType.ArrayDataType
 import com.scalableminds.webknossos.datastore.models.datasource.ElementClass
 
 import java.nio.ByteOrder
@@ -28,7 +28,7 @@ trait DatasetHeader {
 
   lazy val byteOrder: ByteOrder = ByteOrder.BIG_ENDIAN
 
-  lazy val bytesPerElement: Int = bytesPerElementFor(resolvedDataType)
+  lazy val bytesPerElement: Int = ArrayDataType.bytesPerElement(resolvedDataType)
 
   lazy val bytesPerChunk: Int = chunkShape.toList.product * bytesPerElement
 
diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/models/datasource/DataLayer.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/models/datasource/DataLayer.scala
@@ -32,14 +32,6 @@ object DataFormat extends ExtendedEnumeration {
 
 object Category extends ExtendedEnumeration {
   val color, segmentation = Value
-
-  def guessFromElementClass(elementClass: ElementClass.Value): Category.Value =
-    elementClass match {
-      case ElementClass.uint16 => segmentation
-      case ElementClass.uint32 => segmentation
-      case ElementClass.uint64 => segmentation
-      case _                   => color
-    }
 }
 
 object ElementClass extends ExtendedEnumeration {
@@ -186,6 +178,20 @@ object ElementClass extends ExtendedEnumeration {
     case ArrayDataType.i8 => Some(ElementClass.int64)
     case _                => None
   }
+
+  def toArrayDataTypeAndChannel(elementClass: ElementClass.Value): (ArrayDataType, Int) = elementClass match {
+    case ElementClass.uint8  => (ArrayDataType.u1, 1)
+    case ElementClass.uint16 => (ArrayDataType.u2, 1)
+    case ElementClass.uint24 => (ArrayDataType.u1, 3)
+    case ElementClass.uint32 => (ArrayDataType.u4, 1)
+    case ElementClass.uint64 => (ArrayDataType.u8, 1)
+    case ElementClass.float  => (ArrayDataType.f4, 1)
+    case ElementClass.double => (ArrayDataType.f8, 1)
+    case ElementClass.int8   => (ArrayDataType.i1, 1)
+    case ElementClass.int16  => (ArrayDataType.i2, 1)
+    case ElementClass.int32  => (ArrayDataType.i4, 1)
+    case ElementClass.int64  => (ArrayDataType.i8, 1)
+  }
 }
 
 object LayerViewConfiguration {
diff --git a/webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/volume/WKWBucketStreamSink.scala b/webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/volume/WKWBucketStreamSink.scala
@@ -2,15 +2,9 @@ package com.scalableminds.webknossos.tracingstore.tracings.volume
 
 import com.scalableminds.util.geometry.Vec3Int
 import com.scalableminds.util.io.{NamedFunctionStream, NamedStream}
-import com.scalableminds.webknossos.datastore.dataformats.wkw.{
-  ChunkType,
-  VoxelType,
-  WKWDataFormatHelper,
-  WKWFile,
-  WKWHeader
-}
+import com.scalableminds.webknossos.datastore.dataformats.wkw.{ChunkType, WKWDataFormatHelper, WKWFile, WKWHeader}
 import com.scalableminds.webknossos.datastore.models.BucketPosition
-import com.scalableminds.webknossos.datastore.models.datasource.DataLayer
+import com.scalableminds.webknossos.datastore.models.datasource.{DataLayer, ElementClass}
 import com.scalableminds.util.tools.{ByteUtils, Fox, FoxImplicits}
 
 import java.io.DataOutputStream
@@ -24,8 +18,8 @@ class WKWBucketStreamSink(val layer: DataLayer, tracingHasFallbackLayer: Boolean
 
   def apply(bucketStream: Iterator[(BucketPosition, Array[Byte])], mags: Seq[Vec3Int])(
       implicit ec: ExecutionContext): Iterator[NamedStream] = {
-    val (voxelType, numChannels) = VoxelType.fromElementClass(layer.elementClass)
-    val header = WKWHeader(1, DataLayer.bucketLength, ChunkType.LZ4, voxelType, numChannels)
+    val (dataType, numChannels) = ElementClass.toArrayDataTypeAndChannel(layer.elementClass)
+    val header = WKWHeader(1, DataLayer.bucketLength, ChunkType.LZ4, dataType, numChannels)
     bucketStream.flatMap {
       case (bucket, data) =>
         val skipBucket = if (tracingHasFallbackLayer) isRevertedElement(data) else isAllZero(data)
