PySpark：Spark Dataframe - 将 ImageSchema 列转换为 nDArray 作为新列

Posted 2023-04-17

【中文标题】PySpark：Spark Dataframe - 将 ImageSchema 列转换为 nDArray 作为新列

【英文标题】：PySpark: Spark Dataframe - Convert ImageSchema Column to a nDArray as a new column

【发布时间】：2020-03-06 17:16:39

【问题描述】：

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

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

大多数深度学习算法都期望图像是 nDArray。如何在数据框本身中进行这种转换？这是使用udf完成的吗？ 我看到 MMLSpark 中有一些 ndArray 函数集成在 spark 2.4 中，有人用过吗？ 谢谢

【参考方案1】：

我自己想通了，下面的解决方案可以帮助别人。

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")