Bugfix + test

apache · Jan 24, 2018 · 054c1dd · 054c1dd
1 parent 568055d
commit 054c1dd
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/data/mllib/images/multi-channel/BGRA_alpha_60.png b/data/mllib/images/multi-channel/BGRA_alpha_60.png
diff --git a/mllib/src/main/scala/org/apache/spark/ml/image/ImageSchema.scala b/mllib/src/main/scala/org/apache/spark/ml/image/ImageSchema.scala
@@ -169,8 +169,7 @@ object ImageSchema {
         var offset = 0
         for (h <- 0 until height) {
           for (w <- 0 until width) {
-            val color = new Color(img.getRGB(w, h))
-
+            val color = new Color(img.getRGB(w, h), nChannels == 4)
             decoded(offset) = color.getBlue.toByte
             decoded(offset + 1) = color.getGreen.toByte
             decoded(offset + 2) = color.getRed.toByte

diff --git a/mllib/src/test/scala/org/apache/spark/ml/image/ImageSchemaSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/image/ImageSchemaSuite.scala
@@ -103,6 +103,10 @@ class ImageSchemaSuite extends SparkFunSuite with MLlibTestSparkContext {
         -71, -58, -56, -73, -64))),
     "BGRA.png" -> (("CV_8UC4",
       Array[Byte](-128, -128, -8, -1, -128, -128, -8, -1, -128,
-        -128, -8, -1, 127, 127, -9, -1, 127, 127, -9, -1)))
+        -128, -8, -1, 127, 127, -9, -1, 127, 127, -9, -1))),
+    "BGRA_alpha_60.png" -> (("CV_8UC4",
+      Array[Byte](-128, -128, -8, 60, -128, -128, -8, 60, -128,
+        -128, -8, 60, 127, 127, -9, 60, 127, 127, -9, 60)))
   )
+
 }