Skip to content

Instantly share code, notes, and snippets.

@maziyarpanahi

maziyarpanahi/sparknlp-vit-pipeline-base.py

Created August 14, 2022 17:59
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save maziyarpanahi/ae3bf2aca5f0e68435daaf80472f4f5a to your computer and use it in GitHub Desktop.
Save maziyarpanahi/ae3bf2aca5f0e68435daaf80472f4f5a to your computer and use it in GitHub Desktop.
Download ZIP
Raw
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
])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment