sparknlp-vit-pipeline-base.py

from sparknlp.annotator import *
from sparknlp.base import *

from pyspark.ml import Pipeline

imageAssembler = ImageAssembler() \
    .setInputCol("image") \
    .setOutputCol("image_assembler")
    
imageClassifier = ViTForImageClassification \
    .pretrained("image_classifier_vit_base_patch16_224") \
    .setInputCols("image_assembler") \
    .setOutputCol("class") \
    .setBatchSize(8)

pipeline = Pipeline(stages=[
    imageAssembler,
    imageClassifier
])