[SPARK-23205][ML] Update ImageSchema.readImages to correctly set alph…

…a values for four-channel images ## What changes were proposed in this pull request? When parsing raw image data in ImageSchema.decode(), we use a [java.awt.Color](https://docs.oracle.com/javase/7/docs/api/java/awt/Color.html#Color(int)) constructor that sets alpha = 255, even for four-channel images (which may have different alpha values). This PR fixes this issue & adds a unit test to verify correctness of reading four-channel images. ## How was this patch tested? Updates an existing unit test ("readImages pixel values test" in `ImageSchemaSuite`) to also verify correctness when reading a four-channel image. Author: Sid Murching <[email protected]> Closes apache#20389 from smurching/image-schema-bugfix.
sluk3r · Jan 26, 2018 · 7bd46d9 · 7bd46d9
1 parent e57f394
commit 7bd46d9
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 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,12 +169,11 @@ 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), hasAlpha)
             decoded(offset) = color.getBlue.toByte
             decoded(offset + 1) = color.getGreen.toByte
             decoded(offset + 2) = color.getRed.toByte
-            if (nChannels == 4) {
+            if (hasAlpha) {
               decoded(offset + 3) = color.getAlpha.toByte
             }
             offset += nChannels

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
@@ -53,11 +53,11 @@ class ImageSchemaSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(df.count === 1)
 
     df = readImages(imagePath, null, true, -1, false, 1.0, 0)
-    assert(df.count === 9)
+    assert(df.count === 10)
 
     df = readImages(imagePath, null, true, -1, true, 1.0, 0)
     val countTotal = df.count
-    assert(countTotal === 7)
+    assert(countTotal === 8)
 
     df = readImages(imagePath, null, true, -1, true, 0.5, 0)
     // Random number about half of the size of the original dataset
@@ -103,6 +103,9 @@ 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)))
   )
 }