Description
Pretrained VIT model, adapted from Hugging Face and curated to provide scalability and production-readiness using Spark NLP.image_classifier_vit_autotrain_dog_vs_food
is a English model originally trained by abhishek.
Predicted Entities
dog
, food
Live Demo
Open in Colab
Download
Copy S3 URI
How to use
Scala NLU
image_assembler = ImageAssembler () \
. setInputCol ( "image" ) \
. setOutputCol ( "image_assembler" )
imageClassifier = ViTForImageClassification \
. pretrained ( "image_classifier_vit_autotrain_dog_vs_food" , "en" ) \
. setInputCols ( "image_assembler" ) \
. setOutputCol ( "class" )
pipeline = Pipeline ( stages = [
image_assembler ,
imageClassifier ,
])
pipelineModel = pipeline . fit ( imageDF )
pipelineDF = pipelineModel . transform ( imageDF )
val imageAssembler = new ImageAssembler ()
. setInputCol ( "image" )
. setOutputCol ( "image_assembler" )
val imageClassifier = ViTForImageClassification
. pretrained ( "image_classifier_vit_autotrain_dog_vs_food" , "en" )
. setInputCols ( "image_assembler" )
. setOutputCol ( "class" )
val pipeline = new Pipeline (). setStages ( Array ( imageAssembler , imageClassifier ))
val pipelineModel = pipeline . fit ( imageDF )
val pipelineDF = pipelineModel . transform ( imageDF )
import nlu
import requests
response = requests . get ( 'https://raw.githubusercontent.com/JohnSnowLabs/spark-nlp/master/docs/assets/images/hen.JPEG' )
with open ( 'hen.JPEG' , 'wb' ) as f :
f . write ( response . content )
nlu . load ( "en.classify_image.autotrain_dog_vs_food" ). predict ( "hen.JPEG" )
|—|—|
Model Name:
image_classifier_vit_autotrain_dog_vs_food
Compatibility:
Spark NLP 4.1.0+
Input Labels:
[image_assembler]