sparknlp.annotator.cv.vision_encoder_decoder_for_image_captioning#

Classes#

class VisionEncoderDecoderForImageCaptioning(classname='com.johnsnowlabs.nlp.annotators.cv.VisionEncoderDecoderForImageCaptioning', java_model=None)[source]#

VisionEncoderDecoder model that converts images into text captions. It allows for the use of pretrained vision auto-encoding models, such as ViT, BEiT, or DeiT as the encoder, in combination with pretrained language models, like RoBERTa, GPT2, or BERT as the decoder.

Pretrained models can be loaded with pretrained of the companion object:

imageClassifier = VisionEncoderDecoderForImageCaptioning.pretrained() \
    .setInputCols(["image_assembler"]) \
    .setOutputCol("caption")

The default model is "image_captioning_vit_gpt2", if no name is provided.

For available pretrained models please see the Models Hub.

Models from the HuggingFace 🤗 Transformers library are also compatible with Spark NLP 🚀. To see which models are compatible and how to import them see JohnSnowLabs/spark-nlp#5669 and to see more extended examples, see VisionEncoderDecoderTestSpec.

Parameters:

configProtoBytes

ConfigProto from tensorflow, serialized into byte array.

doResize

Whether to resize the input to a certain size

doNormalize

Whether to normalize the input with mean and standard deviation

featureExtractorType

Name of model’s architecture for feature extraction

imageMean

The sequence of means for each channel, to be used when normalizing images

imageStd

The sequence of standard deviations for each channel, to be used when normalizing images

resample

An optional resampling filter. This can be one of PIL.Image.NEAREST, PIL.Image.BILINEAR or PIL.Image.BICUBIC. Only has an effect if do_resize is set to True.

size

Resize the input to the given size. If a tuple is provided, it should be (width, height). If only an integer is provided, then the input will be resized to (size, size). Only has an effect if do_resize is set to True.

doRescale

Whether to rescale the image values by rescaleFactor

rescaleFactor

Factor to scale the image values

minOutputLength

Minimum length of the sequence to be generated

maxOutputLength

Maximum length of output text

doSample

Whether or not to use sampling; use greedy decoding otherwise

temperature

The value used to module the next token probabilities

topK

The number of highest probability vocabulary tokens to keep for top-k-filtering

topP

If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation

repetitionPenalty

The parameter for repetition penalty. 1.0 means no penalty. See this paper for more details

noRepeatNgramSize

If set to int > 0, all ngrams of that size can only occur once

beamSize

The Number of beams for beam search

nReturnSequences

The number of sequences to return from the beam search

Notes

This is a very computationally expensive module especially on larger batch sizes. The use of an accelerator such as GPU is recommended.

Input Annotation types	Output Annotation type
IMAGE	DOCUMENT

Examples

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp.annotator import *
>>> from pyspark.ml import Pipeline
>>> imageDF = spark.read \
...     .format("image") \
...     .option("dropInvalid", value = True) \
...     .load("src/test/resources/image/")
>>> imageAssembler = ImageAssembler() \
...     .setInputCol("image") \
...     .setOutputCol("image_assembler")
>>> imageCaptioning = VisionEncoderDecoderForImageCaptioning \
...     .pretrained() \
...     .setBeamSize(2) \
...     .setDoSample(False) \
...     .setInputCols(["image_assembler"]) \
...     .setOutputCol("caption")
>>> pipeline = Pipeline().setStages([imageAssembler, imageCaptioning])
>>> pipelineDF = pipeline.fit(imageDF).transform(imageDF)
>>> pipelineDF \
...     .selectExpr("reverse(split(image.origin, '/'))[0] as image_name", "caption.result") \
...     .show(truncate = False)
+-----------------+---------------------------------------------------------+
|image_name       |result                                                   |
+-----------------+---------------------------------------------------------+
|palace.JPEG      |[a large room filled with furniture and a large window]  |
|egyptian_cat.jpeg|[a cat laying on a couch next to another cat]            |
|hippopotamus.JPEG|[a brown bear in a body of water]                        |
|hen.JPEG         |[a flock of chickens standing next to each other]        |
|ostrich.JPEG     |[a large bird standing on top of a lush green field]     |
|junco.JPEG       |[a small bird standing on a wet ground]                  |
|bluetick.jpg     |[a small dog standing on a wooden floor]                 |
|chihuahua.jpg    |[a small brown dog wearing a blue sweater]               |
|tractor.JPEG     |[a man is standing in a field with a tractor]            |
|ox.JPEG          |[a large brown cow standing on top of a lush green field]|
+-----------------+---------------------------------------------------------+

setConfigProtoBytes(b)[source]#

Sets configProto from tensorflow, serialized into byte array.

Parameters:

bList[int]

ConfigProto from tensorflow, serialized into byte array

static loadSavedModel(folder, spark_session)[source]#

Loads a locally saved model.

Parameters:

folderstr

Folder of the saved model

spark_sessionpyspark.sql.SparkSession

The current SparkSession

Returns:

VisionEncoderDecoderForImageCaptioning

The restored model

static pretrained(name='image_captioning_vit_gpt2', lang='en', remote_loc=None)[source]#

Downloads and loads a pretrained model.

Parameters:

namestr, optional

Name of the pretrained model, by default “image_captioning_vit_gpt2”

langstr, optional

Language of the pretrained model, by default “en”

remote_locstr, optional

Optional remote address of the resource, by default None. Will use Spark NLPs repositories otherwise.

Returns:

VisionEncoderDecoderForImageCaptioning

The restored model