在 Scala 中使用 K-means 对 Spark 进行图像分割

Posted 2023-03-12

【中文标题】在 Scala 中使用 K-means 对 Spark 进行图像分割

【英文标题】：Image segmentation using K-means for Spark in Scala

【发布时间】：2019-12-14 12:32:25

【问题描述】：

我正在按照教程（来自一本书）使用 Spark 实现用于图像分割的 K-Means 算法。 但是实现是在 Python 中完成的。我认为在 Scala 上实现它会很好。

但我无法通过分割重建图像。

我正在尝试这张来自癌症影像档案 (TCIA) (256x256) 的图片：

这是我的代码：

val spark = SparkSession.builder().appName("mriClass").master("local[*]").getOrCreate()
val mri_healthy_brain_image = "src/main/resources/datasets/clustering/data/mri-images-data/mri-healthy-brain.png"

val image_df = spark.read.format("image").load(mri_healthy_brain_image).select(col("image.*"))
image_df.show
image_df.printSchema
import spark.implicits._

val data = image_df.rdd.collect().map(f => f(5))

val data_array: Array[Byte] = data(0).asInstanceOf[Array[Byte]]

val transposed_df = spark.sparkContext.parallelize(data_array).map(f => Image(f)).toDF

transposed_df.show

val features_col = Array("data")
val vector_assembler = new VectorAssembler()
.setInputCols(features_col)
.setOutputCol("features")

val mri_healthy_brain_df = vector_assembler.transform(transposed_df).select("features")

val k = 5
val kmeans = new KMeans().setK(k).setSeed(12345).setFeaturesCol("features")
val kmeans_model = kmeans.fit(mri_healthy_brain_df)   
val kmeans_centers = kmeans_model.clusterCenters
println("Cluster Centers --------")
for(k <- kmeans_centers)
  println(k)


val mri_healthy_brain_clusters_df = kmeans_model.transform(mri_healthy_brain_df)
.select("features","prediction")

val image_array = mri_healthy_brain_clusters_df.select("prediction").rdd.map(f => f.getAs[Int](0)).collect()

最后，image_array 包含 65536 个位置，每个位置都包含自己的分类

在将图像加载到 Dataframe 时，我认为 spark 会将图像矩阵简单地转换为一维数组，这是 DF 中的二进制类型行。

考虑到这一点，我很容易获取 image_array 并转换为 256x256 图像

我使用地图来预定义分类颜色：

val colors:Map[Int,Int] = Map(
  0 -> 0x717171, //gray
  1 -> 0x0074FF, //light blue
  2 -> 0x95FFDF, //cyan
  3 -> 0xFF3333, //red
  4 -> 0x0058B6, //blue
  5 -> 0xE2CE06, //yellow
  6 -> 0xDB06E2, //pink
  7 -> 0x67C82C, //green
  8 -> 0x8136DC, //purple
  9 -> 0x356F07, //darkgreen
  10 -> 0xE5A812 //orange
  )

并使用此函数生成图像：

def generateImage(img: BufferedImage, image_array: Array[Byte]): BufferedImage = 
// img is the original image
// obtain width and height of image
val w = img.getWidth
val h = img.getHeight

if ( w*h != image_array.size)
  throw new IllegalArgumentException("image array does not fit the provided image");


// create new image of the same size
val out = new BufferedImage(w, h, BufferedImage.TYPE_INT_RGB)

var s = 0
for (x <- 0 until w)
  for (y <- 0 until h)
  out.setRGB(x, y, colors(image_array(s).toInt))
  s +=1
  
out

但我得到的图像是这个：

我可以肯定地说我的集群化管道是正确的，因为它与书中的结果相匹配。

但我不确定 spark 是否在分类后对 Dataframe 上的字节顺序进行排序，可能会损坏结果。

谁能提示我哪里做错了？

提前致谢

【参考方案1】：

我发现了图像数据在 ImageSchema 中的组织方式。图像数据表示为具有维度形状（高度、宽度、nChannels）和类型为 t 的数组值的 3 维数组，由模式字段指定。该数组以行优先顺序存储（在大多数情况下为逐行 BGR）。

由于我对 Open-CV 没有任何经验，并且需要更多时间来理解重建图像的基本原理，我决定使用 Java ImageIO 读取图像，将每个 RGB 信息存储在一个数组中并创建来自它的 DataFrame。

然后我使用之前描述的相同过程，使用 KMeans 分类器，使用带有肿瘤的图像生成预测，并以相同的顺序写入字节来重建图像。

我现在得到的结果是这样的：

你可以在这里找到我的完整代码：

https://github.com/gsjunior86/ScalaMLPratice/blob/master/src/main/scala/br/org/gsj/ml/spark/clustering/kmeans/MriClustering.scala