Spark内置图像数据源初探

概述

在Apache Spark 2.4中引入了一个新的内置数据源, 图像数据源.用户可以通过DataFrame API加载指定目录的中图像文件,生成一个DataFrame对象.通过该DataFrame对象,用户可以对图像数据进行简单的处理,然后使用MLlib进行特定的训练和分类计算.
本文将介绍图像数据源的实现细节和使用方法.

简单使用

先通过一个例子来简单的了解下图像数据源使用方法. 本例设定有一组图像文件存放在阿里云的OSS上, 需要对这组图像加水印,并压缩存储到parquet文件中. 废话不说,先上代码:

  // 为了突出重点,代码简化图像格式相关的处理逻辑def main(args: Array[String]): Unit = {val conf = new SparkConf().setMaster("local[*]")val spark = SparkSession.builder().config(conf).getOrCreate()val imageDF = spark.read.format("image").load("oss://<bucket>/path/to/src/dir")imageDF.select("image.origin", "image.width", "image.height", "image.nChannels", "image.mode", "image.data").map(row => {val origin = row.getAs[String]("origin")val width = row.getAs[Int]("width")val height = row.getAs[Int]("height")val mode = row.getAs[Int]("mode")val nChannels = row.getAs[Int]("nChannels")val data = row.getAs[Array[Byte]]("data")Row(Row(origin, height, width, nChannels, mode,markWithText(width, height, BufferedImage.TYPE_3BYTE_BGR, data, "EMR")))}).write.format("parquet").save("oss://<bucket>/path/to/dst/dir")}def markWithText(width: Int, height: Int, imageType: Int, data: Array[Byte], text: String): Array[Byte] = {val image = new BufferedImage(width, height, imageType)val raster = image.getData.asInstanceOf[WritableRaster]val pixels = data.map(_.toInt)raster.setPixels(0, 0, width, height, pixels)image.setData(raster)val buffImg = new BufferedImage(width, height, imageType)val g = buffImg.createGraphicsg.drawImage(image, 0, 0, null)g.setColor(Color.red)g.setFont(new Font("宋体", Font.BOLD, 30))g.drawString(text, width/2, height/2)g.dispose()val buffer = new ByteArrayOutputStreamImageIO.write(buffImg, "JPG", buffer)buffer.toByteArray}

从生成的parquet文件中抽取一条图像二进制数据,保存为本地jpg,效果如下:

图1 左图为原始图像,右图为处理后的图像

你可能注意到两个图像到颜色并不相同,这是因为Spark的图像数据将图像解码为BGR顺序的数据,而示例程序在保存的时候,没有处理这个变换,导致颜色出现了反差.

实现初窥

下面我们深入到spark源码中来看一下实现细节.Apache Spark内置图像数据源的实现代码在spark-mllib这个模块中.主要包括两个类:

org.apache.spark.ml.image.ImageSchema
org.apache.spark.ml.source.image.ImageFileFormat

其中,ImageSchema定义了图像文件加载为DataFrame的Row的格式和解码方法.ImageFileFormat提供了面向存储层的读写接口.

格式定义

一个图像文件被加载为DataFrame后,对应的如下:

  val columnSchema = StructType(StructField("origin", StringType, true) ::StructField("height", IntegerType, false) ::StructField("width", IntegerType, false) ::StructField("nChannels", IntegerType, false) ::// OpenCV-compatible type: CV_8UC3 in most casesStructField("mode", IntegerType, false) ::// Bytes in OpenCV-compatible order: row-wise BGR in most casesStructField("data", BinaryType, false) :: Nil)val imageFields: Array[String] = columnSchema.fieldNamesval imageSchema = StructType(StructField("image", columnSchema, true) :: Nil)

如果将该DataFrame打印出来,可以得到如下形式的表:

+--------------------+-----------+------------+---------------+----------+-------------------+
|image.origin        |image.width|image.height|image.nChannels|image.mode|image.data         |
+--------------------+-----------+------------+---------------+----------+-------------------+
|oss://.../dir/1.jpg |600        |343         |3              |16        |55 45 21 56  ...   |
+--------------------+-----------+------------+---------------+----------+-------------------+

其中:

origin: 原始图像文件的路径
width: 图像的宽度,单位像素
height: 图像的高度,单位像素
nChannels: 图像的通道数, 如常见的RGB位图为通道数为3
mode: 像素矩阵(data)中元素的数值类型和通道顺序, 与OpenCV的类型兼容
data: 解码后的像素矩阵

提示: 关于图像的基础支持,可以参考如下文档: Image file reading and writing

加载和解码

图像文件通过ImageFileFormat加载为一个Row对象.

// 文件: ImageFileFormat.scala
// 为了简化说明起见,代码有删减和改动
private[image] class ImageFileFormat extends FileFormat with DataSourceRegister {......override def prepareWrite(sparkSession: SparkSession,job: Job,options: Map[String, String],dataSchema: StructType): OutputWriterFactory = {throw new UnsupportedOperationException("Write is not supported for image data source")}override protected def buildReader(sparkSession: SparkSession,dataSchema: StructType,partitionSchema: StructType,requiredSchema: StructType,filters: Seq[Filter],options: Map[String, String],hadoopConf: Configuration): (PartitionedFile) => Iterator[InternalRow] = {    ......(file: PartitionedFile) => {......val path = new Path(origin)val stream = fs.open(path)val bytes = ByteStreams.toByteArray(stream)val resultOpt = ImageSchema.decode(origin, bytes) // <-- 解码 val filteredResult = Iterator(resultOpt.getOrElse(ImageSchema.invalidImageRow(origin)))......val converter = RowEncoder(requiredSchema)filteredResult.map(row => converter.toRow(row))......}}}
}

从上可以看出:

当前的图像数据源实现并不支持保存操作;
图像数据的解码工作在ImageSchema中完成.

下面来看一下具体的解码过程:

// 文件: ImageSchema.scala
// 为了简化说明起见,代码有删减和改动
private[spark] def decode(origin: String, bytes: Array[Byte]): Option[Row] = {// 使用ImageIO加载原始图像数据val img = ImageIO.read(new ByteArrayInputStream(bytes))if (img != null) {// 获取图像的基本属性val isGray = img.getColorModel.getColorSpace.getType == ColorSpace.TYPE_GRAYval hasAlpha = img.getColorModel.hasAlphaval height = img.getHeightval width = img.getWidth// ImageIO::ImageType -> OpenCV Typeval (nChannels, mode) = if (isGray) {(1, ocvTypes("CV_8UC1"))} else if (hasAlpha) {(4, ocvTypes("CV_8UC4"))} else {(3, ocvTypes("CV_8UC3"))}// 解码val imageSize = height * width * nChannels// 用于存储解码后的像素矩阵val decoded = Array.ofDim[Byte](imageSize)if (isGray) {// 处理单通道图像...} else {// 处理多通道图像var offset = 0for (h <- 0 until height) {for (w <- 0 until width) {val color = new Color(img.getRGB(w, h), hasAlpha)// 解码后的通道顺序为BGR(A)decoded(offset) = color.getBlue.toBytedecoded(offset + 1) = color.getGreen.toBytedecoded(offset + 2) = color.getRed.toByteif (hasAlpha) {decoded(offset + 3) = color.getAlpha.toByte}offset += nChannels}}}// 转换为一行数据Some(Row(Row(origin, height, width, nChannels, mode, decoded)))}}

从上可以看出:

本数据源在实现上使用javax的ImageIO库实现各类格式的图像文件的解码.ImageIO虽然是一个十分强大和专业的java图像处理库,但是和更专业的CV库(如OpenCV)比起来,性能上和功能上差距还是非常大的;
解码后的图像通道顺序和像素数值类型是固定的, 顺序固定为BGR(A), 像素数值类型为8U;
最多支持4个通道,因此像多光谱遥感图像这类可能包含数十个波段信息的图像就无法支持了;
解码后输出的信息仅包含基本的长宽、通道数和模式等字段,如果需要获取更为详细元数据,如exif,GPS坐标等就爱莫能助了;
数据源在生成DataFrame时执行了图像的解码操作,并且解码后的数据存储在Java堆内内存中.这在实际项目应该是一个比较粗放的实现方式,会占用大量的资源,包括内存和带宽(如果发生shuffle的话,可以考虑参考同一个图像文件保存为BMP和JPG的大小差别).

编码和存储

从上分析可以看出,当前图像数据源并不支持对处理后的像素矩阵进行编码并保存为指定格式的图像文件.

图像处理能力

当前版本Apache Spark并没有提供面向图像数据的UDF,图像数据的处理需要借助ImageIO库或其他更专业的CV库.

小结

当前Apache Spark的内置图像数据源可以较为方便的加载图像文件进行分析.不过,当前的实现还十分简陋,性能和资源消耗应该都不会太乐观.并且,当前版本仅提供了图像数据的加载能力,并没有提供常用处理算法的封装和实现,也不能很好的支持更为专业的CV垂直领域的分析业务.当然,这和图像数据源在Spark中的定位有关(将图像数据作为输入用于训练DL模型,这类任务对图像的处理本身要求并不多).如果希望使用Spark框架完成更实际的图像处理任务,还有很多工作要做,比如: