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

Question

我正在按照教程（来自一本书）使用 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 上的字节顺序进行排序，可能会破坏结果。

谁能告诉我哪里做错了？

提前致谢

Answer 1

我找到了图像数据在 ImageSchema 中的组织方式。图像数据表示为一个 3 维数组，其维度形状（高度、宽度、nChannels）和模式字段指定的类型 t 的数组值。该数组以行优先顺序存储（大多数情况下是按行 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

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

Image segmentation using K-means for Spark in Scala

scala

apache-spark

machine-learning

k-means

apache-spark-ml