Update ImageUtils to support resizing one, three, or four channel images

src/main/scala/com/databricks/sparkdl/DeepImageFeaturizer.scala

@@ -115,8 +115,9 @@ class DeepImageFeaturizer(override val uid: String) extends Transformer with Def
     val height = model.height
     val width = model.width
 
-    val resizeUdf = udf((image: Row) => ImageUtils.resizeImage(height, width, 3, image,
-      DeepImageFeaturizer.scaleHints(getScaleHint)), imSchema)
+    val resizeUdf = udf((image: Row) => { ImageUtils.resizeImage(height, width, 3, image,
+          DeepImageFeaturizer.scaleHints(getScaleHint))
+      }, imSchema)
 
     val imageDF = dataFrame
       .withColumn(RESIZED_IMAGE_COL, resizeUdf(col(getInputCol)))

src/main/scala/com/databricks/sparkdl/ImageUtils.scala

@@ -16,6 +16,7 @@
 
 package com.databricks.sparkdl
 
+import java.awt.color.ColorSpace
 import java.awt.image.BufferedImage
 import java.awt.{Color, Image}
 
@@ -25,8 +26,8 @@ import org.apache.spark.sql.Row
 private[sparkdl] object ImageUtils {
 
   /**
-   * Takes a Row image (spImage) and returns a Java BufferedImage. Currently supports 1 & 3
-   * channel images. If the image has 3 channels, we assume the channels are in BGR order.
+   * Takes a Row image (spImage) and returns a Java BufferedImage. Currently supports 1, 3, & 4
+   * channel images. If the image has 3 or 4 channels, we assume the channels are in BGR(A) order.
    *
    * @param rowImage Image in spark.ml.image format.
    * @return Java BGR BufferedImage.
@@ -42,10 +43,15 @@ private[sparkdl] object ImageUtils {
           | image of size ($height, $width, $channels).
        """.stripMargin
     )
-    val image = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR)
+
+    val image = channels match {
+      case 1 => new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY)
+      case 3 => new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR)
+      case 4 => new BufferedImage(width, height, BufferedImage.TYPE_4BYTE_ABGR)
+    }
 
     var offset, h = 0
-    var r, g, b: Byte = 0
+    var r, g, b, a: Byte = 0
     while (h < height) {
       var w = 0
       while (w < width) {
@@ -58,11 +64,20 @@ private[sparkdl] object ImageUtils {
             b = imageData(offset)
             g = imageData(offset + 1)
             r = imageData(offset + 2)
+          case 4 =>
+            b = imageData(offset)
+            g = imageData(offset + 1)
+            r = imageData(offset + 2)
+            a = imageData(offset + 3)
           case _ =>
-            require(false, s"`Channels` must be 1 or 3, got $channels.")
+            require(false, s"`Channels` must be 1, 3, or 4, got $channels.")
         }
 
-        val color = new Color(r & 0xff, g & 0xff, b & 0xff)
+        val color = if (channels < 4) {
+          new Color(r & 0xff, g & 0xff, b & 0xff)
+        } else {
+          new Color(r & 0xff, g & 0xff, b & 0xff, a & 0xff)
+        }
         image.setRGB(w, h, color.getRGB)
         offset += channels
         w += 1
@@ -72,15 +87,40 @@ private[sparkdl] object ImageUtils {
     image
   }
 
+  /** Returns the number of channels in the passed-in buffered image. */
+  private def getNumChannels(img: BufferedImage): Int = {
+    val isGray = img.getColorModel.getColorSpace.getType == ColorSpace.TYPE_GRAY
+    val hasAlpha = img.getColorModel.hasAlpha
+    if (isGray) {
+      1
+    } else if (hasAlpha) {
+      4
+    } else {
+      3
+    }
+  }
+
+  /** Returns the OCV type (int) of the passed-in image */
+  private def getOCVType(img: BufferedImage): Int = {
+    val isGray = img.getColorModel.getColorSpace.getType == ColorSpace.TYPE_GRAY
+    val hasAlpha = img.getColorModel.hasAlpha
+    if (isGray) {
+      ImageSchema.ocvTypes("CV_8UC1")
+    } else if (hasAlpha) {
+      ImageSchema.ocvTypes("CV_8UC4")
+    } else {
+      ImageSchema.ocvTypes("CV_8UC3")
+    }
+  }
 
   /**
    * Takes a Java BufferedImage and returns a Row Image (spImage).
    *
    * @param image Java BufferedImage.
-   * @return Row image in spark.ml.image format with 3 channels in BGR order.
+   * @return Row image in spark.ml.image format with channels in BGR(A) order.
    */
   private[sparkdl] def spImageFromBufferedImage(image: BufferedImage, origin: String = null): Row = {
-    val channels = 3
+    val channels = getNumChannels(image)
     val height = image.getHeight
     val width = image.getWidth
 
@@ -89,27 +129,37 @@ private[sparkdl] object ImageUtils {
     while (h < height) {
       var w = 0
       while (w < width) {
-        val color = new Color(image.getRGB(w, h))
-        decoded(offset) = color.getBlue.toByte
-        decoded(offset + 1) = color.getGreen.toByte
-        decoded(offset + 2) = color.getRed.toByte
+        val color = new Color(image.getRGB(w, h), image.getColorModel.hasAlpha)
+        channels match {
+          case 1 =>
+            decoded(offset) = color.getBlue.toByte
+          case 3 =>
+            decoded(offset) = color.getBlue.toByte
+            decoded(offset + 1) = color.getGreen.toByte
+            decoded(offset + 2) = color.getRed.toByte
+          case 4 =>
+            decoded(offset) = color.getBlue.toByte
+            decoded(offset + 1) = color.getGreen.toByte
+            decoded(offset + 2) = color.getRed.toByte
+            decoded(offset + 3) = color.getAlpha.toByte
+        }
         offset += channels
         w += 1
       }
       h += 1
     }
-    Row(origin, height, width, channels, ImageSchema.ocvTypes("CV_8UC3"), decoded)
+    Row(origin, height, width, channels, getOCVType(image), decoded)
+
   }
 
   /**
-   * Resizes an image and returns it as an Array[Byte]. Only 1 and 3 channel inputs, where each
+   * Resizes an image and returns it as an Array[Byte]. Only 1, 3, and 4 channel inputs, where each
    * channel is a single Byte, are currently supported. Only BGR channel order is supported but
    * this might work for other channel orders.
    *
    * @param tgtHeight   desired height of output image.
    * @param tgtWidth    desired width of output image.
-   * @param tgtChannels number of channels of output image (must be 3), may be used later to
-   *                    support more channels.
+   * @param tgtChannels number of channels in output image.
    * @param spImage     image to resize.
    * @param scaleHint   hint which algorhitm to use, see java.awt.Image#SCALE_SCALE_AREA_AVERAGING
    * @return resized image, if the input was BGR or 1 channel, the output will be BGR.
@@ -120,8 +170,6 @@ private[sparkdl] object ImageUtils {
     tgtChannels: Int,
     spImage: Row,
     scaleHint: Int = Image.SCALE_AREA_AVERAGING): Row = {
-    require(tgtChannels == 3, s"`tgtChannels` was set to $tgtChannels, must be 3.")
-
     val height = ImageSchema.getHeight(spImage)
     val width = ImageSchema.getWidth(spImage)
     val nChannels = ImageSchema.getNChannels(spImage)
@@ -130,14 +178,19 @@ private[sparkdl] object ImageUtils {
       spImage
     } else {
       val srcImg = spImageToBufferedImage(spImage)
-      val tgtImg = new BufferedImage(tgtWidth, tgtHeight, BufferedImage.TYPE_3BYTE_BGR)
+      val tgtImgType = tgtChannels match {
+        case 1 => BufferedImage.TYPE_BYTE_GRAY
+        case 3 => BufferedImage.TYPE_3BYTE_BGR
+        case 4 => BufferedImage.TYPE_4BYTE_ABGR
+      }
+      val tgtImg = new BufferedImage(tgtWidth, tgtHeight, tgtImgType)
       // scaledImg is a java.awt.Image which supports drawing but not pixel lookup by index.
       val scaledImg = srcImg.getScaledInstance(tgtWidth, tgtHeight, scaleHint)
       // Draw scaledImage onto resized (usually smaller) tgtImg so extract individual pixel values.
       val graphic = tgtImg.createGraphics()
       graphic.drawImage(scaledImg, 0, 0, null)
       graphic.dispose()
-      spImageFromBufferedImage(tgtImg, origin=ImageSchema.getOrigin(spImage))
+      spImageFromBufferedImage(tgtImg, origin = ImageSchema.getOrigin(spImage))
     }
   }
 }

src/test/scala/com/databricks/sparkdl/ImageUtilsSuite.scala

@@ -22,54 +22,57 @@ import javax.imageio.ImageIO
 
 import scala.util.Random
 
+import org.scalatest.FunSuite
+
 import org.apache.spark.ml.image.ImageSchema
 import org.apache.spark.sql.Row
 
-import org.scalatest.FunSuite
-
 object ImageUtilsSuite {
-  val biggerImage: Row = {
-    val biggerFile = getClass.getResource("/sparkdl/test-image-collection/00081101.jpg").getFile
-    val imageBuffer = ImageIO.read(new File(biggerFile))
-    ImageUtils.spImageFromBufferedImage(imageBuffer)
-  }
 
-  val smallerImage: Row = {
-    val smallerFile = getClass.getResource("/sparkdl/00081101-small-version.png").getFile
-    val imageBuffer = ImageIO.read(new File(smallerFile))
+  /** Read image data into a BufferedImage, then use our utility method to convert to a row image */
+  def getImageRow(resourcePath: String): Row = {
+    val resourceUrl = getClass.getResource(resourcePath).getFile
+    val imageBuffer = ImageIO.read(new File(resourceUrl))
     ImageUtils.spImageFromBufferedImage(imageBuffer)
   }
 
-
+  def smallerImage: Row = getImageRow("/sparkdl/00081101-small-version.png")
+  def biggerImage: Row = getImageRow("/sparkdl/test-image-collection/00081101.jpg")
 }
 
 class ImageUtilsSuite extends FunSuite {
   // We want to make sure to test ImageUtils in headless mode to ensure it'll work on all systems.
   assert(System.getProperty("java.awt.headless") === "true")
-
-  import ImageUtilsSuite._
-
   test("Test spImage resize.") {
-    val tgtHeight: Int = ImageSchema.getHeight(smallerImage)
-    val tgtWidth: Int = ImageSchema.getWidth(smallerImage)
-    val tgtChannels: Int = ImageSchema.getNChannels(smallerImage)
+    def getImagePath(imageSize: String, numChannels: Int): String = {
+      s"/sparkdl/test-image-collection/${numChannels}_channels/$imageSize.png"
+    }
+    for (channels <- Seq(1, 3, 4)) {
+      val smallerImage = ImageUtilsSuite.getImageRow(getImagePath("small", channels))
+      val biggerImage = ImageUtilsSuite.getImageRow(getImagePath("big", channels))
 
-    val testImage = ImageUtils.resizeImage(tgtHeight, tgtWidth, tgtChannels, biggerImage)
-    assert(testImage === smallerImage, "Resizing image did not produce expected smaller image.")
+      val tgtHeight: Int = ImageSchema.getHeight(smallerImage)
+      val tgtWidth: Int = ImageSchema.getWidth(smallerImage)
+      val tgtChannels: Int = ImageSchema.getNChannels(smallerImage)
+
+      val testImage = ImageUtils.resizeImage(tgtHeight, tgtWidth, tgtChannels, biggerImage)
+      assert(testImage === smallerImage, "Resizing image did not produce expected smaller image.")
+    }
   }
 
   test ("Test Row image -> BufferedImage -> Row image") {
     val height = 200
     val width = 100
-    val channels = 3
-
-    val rand = new Random(971)
-    val imageData = Array.ofDim[Byte](height * width * channels)
-    rand.nextBytes(imageData)
-    val spImage = Row(null, height, width, channels, ImageSchema.ocvTypes("CV_8UC3"), imageData)
-    val bufferedImage = ImageUtils.spImageToBufferedImage(spImage)
-    val testImage = ImageUtils.spImageFromBufferedImage(bufferedImage)
-    assert(spImage === testImage, "Image changed during conversion.")
+    for (channels <- Seq(3, 4)) {
+      val rand = new Random(971)
+      val imageData = Array.ofDim[Byte](height * width * channels)
+      rand.nextBytes(imageData)
+      val ocvType = s"CV_8UC$channels"
+      val spImage = Row(null, height, width, channels, ImageSchema.ocvTypes(ocvType), imageData)
+      val bufferedImage = ImageUtils.spImageToBufferedImage(spImage)
+      val testImage = ImageUtils.spImageFromBufferedImage(bufferedImage)
+      assert(spImage === testImage, s"Image changed during conversion")
+    }
   }
 
   test("Simple BufferedImage from Row Image") {