PySpark: Spark Dataframe - конвертировать столбец ImageSchema в nDArray как новый столбец

Question

Я читаю изображения (* .jpg) в папке с искровым фреймом данных, используя ImageSchema.

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

• Большинство алгоритмов DeepLearning ожидают, что изображения будут nDArray. Как сделать это преобразование в самом фрейме данных? Это должно быть сделано с помощью UDF?
  
• Я вижу, что есть некоторые функции ndArray в MMLSpark, которые были интегрированы в spark 2.4 , кто-нибудь использовал это?
  Спасибо

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