Change dimensions order in the image shape from WHC to HWC

Kotlin · Apr 24, 2023 · 1548e0e · 1548e0e
1 parent a8b578f
commit 1548e0e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/impl/src/jvmMain/kotlin/org/jetbrains/kotlinx/dl/impl/preprocessing/image/ImageUtil.kt b/impl/src/jvmMain/kotlin/org/jetbrains/kotlinx/dl/impl/preprocessing/image/ImageUtil.kt
@@ -25,5 +25,5 @@ internal fun BufferedImage.copy(): BufferedImage {
 }
 
 internal fun BufferedImage.getShape(): TensorShape {
-    return TensorShape(width.toLong(), height.toLong(), colorModel.numComponents.toLong())
+    return TensorShape(height.toLong(), width.toLong(), colorModel.numComponents.toLong())
 }
diff --git a/...jvmTest/kotlin/org/jetbrains/kotlinx/dl/impl/preprocessing/PreprocessingFinalShapeTest.kt b/...jvmTest/kotlin/org/jetbrains/kotlinx/dl/impl/preprocessing/PreprocessingFinalShapeTest.kt
@@ -127,7 +127,7 @@ class PreprocessingFinalShapeTest {
         val image = BufferedImage(10, 20, BufferedImage.TYPE_3BYTE_BGR)
         val (_, actualShape) = preprocess.apply(image)
 
-        assertEquals(actualShape, preprocess.getOutputShape(TensorShape(10, 20, 1)))
+        assertEquals(actualShape, preprocess.getOutputShape(TensorShape(20, 10, 1)))
     }
 
     @Test
@@ -139,6 +139,6 @@ class PreprocessingFinalShapeTest {
         val image = BufferedImage(10, 20, BufferedImage.TYPE_3BYTE_BGR)
         val (_, actualShape) = preprocess.apply(image)
 
-        assertEquals(actualShape, preprocess.getOutputShape(TensorShape(10, 20, 3)))
+        assertEquals(actualShape, preprocess.getOutputShape(TensorShape(20, 10, 3)))
     }
 }
diff --git a/.../src/jvmTest/kotlin/org/jetbrains/kotlinx/dl/impl/preprocessing/PreprocessingImageTest.kt b/.../src/jvmTest/kotlin/org/jetbrains/kotlinx/dl/impl/preprocessing/PreprocessingImageTest.kt
@@ -105,7 +105,7 @@ class PreprocessingImageTest {
         inputImage.setRGB(1, 1, Color.RED.rgb)
         val (imageFloats, tensorShape) = preprocess.apply(inputImage)
 
-        Assertions.assertEquals(TensorShape(9, 5, 3), tensorShape)
+        Assertions.assertEquals(TensorShape(5, 9, 3), tensorShape)
 
         val expectedImage = FloatArray(tensorShape.numElements().toInt()) { Color.GRAY.red / 255f }
         expectedImage.setRGB(3, 1, Color.BLUE, tensorShape, ColorMode.BGR)
@@ -198,7 +198,7 @@ class PreprocessingImageTest {
                 }
             }
             for (i in colorComponents.indices) {
-                set3D(y, x, i, tensorShape[0].toInt(), colorMode.channels, colorComponents[i])
+                set3D(y, x, i, tensorShape[1].toInt(), colorMode.channels, colorComponents[i])
             }
         }
     }