PySpark：Spark数据框-将ImageSchema列转换为nDArray作为新列

Posted 2021-04-01

我正在使用ImageSchema读取文件夹中关于spark数据帧的图像（* .jpg）。

image_df = spark.read.format("image").load("/mnt/train/*", inferschema=True)

• [大多数DeepLearning算法期望图像为nDArray。如何在数据框本身中进行此转换？要使用udf完成吗？
• 我看到spark 2.4中集成了MMLSpark中的ndArray函数，有人用过吗？谢谢

答案

我自己弄清楚了，下面的解决方案可以帮助某人。

from pyspark.sql.types import ArrayType,IntegerType
from pyspark.sql.functions import regexp_replace
import numpy as np

def to_np_array(x):    
  height = 200
  width  = 200
  nChannels = 3
  return np.reshape(x, (height,width,nChannels)).tolist()
spark_to_np_array = udf(to_np_array, ArrayType(ArrayType(ArrayType(IntegerType()))))

imagesdf = spark.read.format("image").load("/mnt/images/*", inferschema=True)
imagesdf = imagesdf.withColumn("FileName", regexp_replace('image.origin', 'dbfs:/mnt/images/', '')) 
imagesdf = imagesdf.withColumn("ImageArray", spark_to_np_array(imagesdf["image.data"])).select("FileName","ImageArray")