!pip install --upgrade pip --quiet
!pip install transformers==4.30.2 --quiet
!pip install timm==0.4.12 --quiet
!pip install ftfy==6.1.1 --quiet
!pip install spacy==3.4.1 --quiet
!pip install fairscale==0.4.4 --quiet
!pip install apache_beam[gcp]>=2.48.0

# To use the newly installed versions, restart the runtime.
exit()

import requests
import os
import urllib
import json
import io
from io import BytesIO
from typing import Sequence
from typing import Iterator
from typing import Iterable
from typing import Tuple
from typing import Optional
from typing import Dict
from typing import List
from typing import Any

import apache_beam as beam
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.ml.inference.base import KeyedModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor
from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor
from transformers import CLIPProcessor
from transformers import CLIPTokenizer
from transformers import CLIPModel
from transformers import CLIPConfig
from transformers import CLIPFeatureExtractor
import torch
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

 git lfs install
 git clone https://huggingface.co/openai/clip-vit-base-patch32

Git LFS initialized.
Cloning into 'clip-vit-base-patch32'...
remote: Enumerating objects: 51, done.
remote: Counting objects: 100% (6/6), done.
remote: Compressing objects: 100% (6/6), done.
remote: Total 51 (delta 1), reused 0 (delta 0), pack-reused 45
Unpacking objects: 100% (51/51), done.
Filtering content: 100% (3/3), 1.69 GiB | 66.58 MiB/s, done.

# CLIP model and component configs paths
clip_feature_extractor_config_path = '/content/clip-vit-base-patch32/preprocessor_config.json'
clip_tokenizer_vocab_config_path = '/content/clip-vit-base-patch32/vocab.json'
clip_merges_config_path = '/content/clip-vit-base-patch32/merges.txt'
clip_model_config_path = '/content/clip-vit-base-patch32/config.json'
clip_state_dict_path = '/content/clip-vit-base-patch32/pytorch_model.bin'

!git clone https://github.com/salesforce/BLIP
%cd /content/BLIP

Cloning into 'BLIP'...
remote: Enumerating objects: 274, done.
remote: Counting objects: 100% (156/156), done.
remote: Compressing objects: 100% (39/39), done.
remote: Total 274 (delta 130), reused 117 (delta 117), pack-reused 118
Receiving objects: 100% (274/274), 7.04 MiB | 13.40 MiB/s, done.
Resolving deltas: 100% (150/150), done.
/content/BLIP

from BLIP.models.blip import blip_decoder

!gdown 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_base_caption.pth'
# The blip model is saved as a checkpoint, load it and save it as a state dict since RunInference required
# a state dict for model instantiation
blip_state_dict_path = '/content/BLIP/blip_state_dict.pth'
torch.save(torch.load('/content/BLIP/model*_base_caption.pth')['model'], blip_state_dict_path)

Downloading...
From: https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_base_caption.pth
To: /content/BLIP/model*_base_caption.pth
100% 896M/896M [00:04<00:00, 198MB/s]

class ReadImagesFromUrl(beam.DoFn):
  """
  Read an image from a given URL and return a tuple of the images_url
  and image data.
  """
  def process(self, element: str) -> Tuple[str, Image.Image]:
    response = requests.get(element)
    image = Image.open(BytesIO(response.content)).convert('RGB')
    return [(element, image)]


class FormatCaptions(beam.DoFn):
  """
  Print the image name and its most relevant captions after CLIP ranking.
  """
  def __init__(self, number_of_top_captions: int):
    self._number_of_top_captions = number_of_top_captions

  def process(self, element: Tuple[str, List[str]]):
    image_url, caption_list = element
    caption_list = caption_list[:self._number_of_top_captions]
    img_name = os.path.basename(image_url).rsplit('.')[0]
    print(f'Image: {img_name}')
    print(f'\tTop {self._number_of_top_captions} captions ranked by CLIP:')
    for caption_rank, caption_prob_pair in enumerate(caption_list):
      print(f'\t\t{caption_rank+1}: {caption_prob_pair[0]}. (Caption probability: {caption_prob_pair[1]:.2f})')
    print('\n')

class PreprocessBLIPInput(beam.DoFn):

  """
  Process the raw image input to a format suitable for BLIP inference. The processed
  images are duplicated to the number of desired captions per image.

  Preprocessing transformation taken from:
  https://github.com/salesforce/BLIP/blob/d10be550b2974e17ea72e74edc7948c9e5eab884/predict.py
  """

  def __init__(self, captions_per_image: int):
    self._captions_per_image = captions_per_image

  def setup(self):

    # Initialize the image transformer.
    self._transform = transforms.Compose([
      transforms.Resize((384, 384),interpolation=InterpolationMode.BICUBIC),
      transforms.ToTensor(),
      transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
    ])

  def process(self, element):
    image_url, image = element
    # The following lines provide a workaround to turn off BatchElements.
    preprocessed_img = self._transform(image).unsqueeze(0)
    preprocessed_img = preprocessed_img.repeat(self._captions_per_image, 1, 1, 1)
    # Parse the processed input to a dictionary to a format suitable for RunInference.
    preprocessed_dict = {'inputs': preprocessed_img}

    return [(image_url, preprocessed_dict)]

class PostprocessBLIPOutput(beam.DoFn):
  """
  Process the PredictionResult to get the generated image captions
  """
  def process(self, element : Tuple[str, Iterable[PredictionResult]]):
    image_url, prediction = element

    return [(image_url, prediction.inference)]

class PreprocessCLIPInput(beam.DoFn):

  """
  Process the image-caption pair to a format suitable for CLIP inference.

  After grouping the raw images with the generated captions, we need to
  preprocess them before passing them to the ranking stage (CLIP model).
  """

  def __init__(self,
               feature_extractor_config_path: str,
               tokenizer_vocab_config_path: str,
               merges_file_config_path: str):

    self._feature_extractor_config_path = feature_extractor_config_path
    self._tokenizer_vocab_config_path = tokenizer_vocab_config_path
    self._merges_file_config_path = merges_file_config_path


  def setup(self):

    # Initialize the CLIP feature extractor.
    feature_extractor_config = CLIPConfig.from_pretrained(self._feature_extractor_config_path)
    feature_extractor = CLIPFeatureExtractor(feature_extractor_config)

    # Initialize the CLIP tokenizer.
    tokenizer = CLIPTokenizer(self._tokenizer_vocab_config_path,
                              self._merges_file_config_path)

    # Initialize the CLIP processor used to process the image-caption pair.
    self._processor = CLIPProcessor(feature_extractor=feature_extractor,
                                    tokenizer=tokenizer)

  def process(self, element: Tuple[str, Dict[str, List[Any]]]):

    image_url, image_captions_pair = element
    # Unpack the image and captions after grouping them with 'CoGroupByKey()'.
    image = image_captions_pair['image'][0]
    captions = image_captions_pair['captions'][0]
    preprocessed_clip_input = self._processor(images = image,
                                              text = captions,
                                              return_tensors="pt",
                                              padding=True)

    image_url_caption_pair = (image_url, captions)
    return [(image_url_caption_pair, preprocessed_clip_input)]


class RankCLIPOutput(beam.DoFn):
  """
  Process the output of CLIP to get the captions sorted by ranking order.

  The logits are the output of the CLIP model. Here, we apply a softmax activation
  function to the logits to get the probabilistic distribution of the relevance
  of each caption to the target image. After that, we sort the captions in descending
  order with respect to the probabilities as a caption-probability pair.
  """

  def process(self, element : Tuple[Tuple[str, List[str]], Iterable[PredictionResult]]):
    (image_url, captions), prediction = element
    prediction_results = prediction.inference
    prediction_probs = prediction_results.softmax(dim=-1).cpu().detach().numpy()
    ranking = np.argsort(-prediction_probs)
    sorted_caption_prob_pair = [(captions[idx], prediction_probs[idx]) for idx in ranking]

    return [(image_url, sorted_caption_prob_pair)]

MAX_CAPTION_LENGTH = 80
MIN_CAPTION_LENGTH = 10
# Increasing Beam search might improve the quality of the captions,
# but also results in more compute time
NUM_BEAMS = 1

def blip_keyed_tensor_inference_fn(
    batch: Sequence[Dict[str, torch.Tensor]],
    model: torch.nn.Module,
    device: str,
    inference_args: Optional[Dict[str, Any]] = None,
    model_id: Optional[str] = None,
) -> Iterable[PredictionResult]:
  # By default, Beam batches inputs for bulk inference and calls model(batch)
  # Since we want to call model.generate on a single unbatched input (BLIP/CLIP
  # don't handle batched inputs), we define a custom inference function.
  captions = model.generate(batch[0]['inputs'],
                            sample=True,
                            num_beams=NUM_BEAMS,
                            max_length=MAX_CAPTION_LENGTH,
                            min_length=MIN_CAPTION_LENGTH)
  return [PredictionResult(batch[0], captions, model_id)]


BLIP_model_handler = PytorchModelHandlerKeyedTensor(
    state_dict_path=blip_state_dict_path,
    model_class=blip_decoder,
    inference_fn=blip_keyed_tensor_inference_fn,
    max_batch_size=1)

BLIP_keyed_model_handler = KeyedModelHandler(BLIP_model_handler)

def clip_keyed_tensor_inference_fn(
    batch: Sequence[Dict[str, torch.Tensor]],
    model: torch.nn.Module,
    device: str,
    inference_args: Optional[Dict[str, Any]] = None,
    model_id: Optional[str] = None,
) -> Iterable[PredictionResult]:
  # By default, Beam batches inputs for bulk inference and calls model(batch)
  # Since we want to call model on a single unbatched input (BLIP/CLIP don't
  # handle batched inputs), we define a custom inference function.
  output = model(**batch[0], **inference_args)
  return [PredictionResult(batch[0], output.logits_per_image[0], model_id)]


CLIP_model_handler = PytorchModelHandlerKeyedTensor(
    state_dict_path=clip_state_dict_path,
    model_class=CLIPModel,
    model_params={'config': CLIPConfig.from_pretrained(clip_model_config_path)},
    inference_fn=clip_keyed_tensor_inference_fn,
    max_batch_size=1)

CLIP_keyed_model_handler = KeyedModelHandler(CLIP_model_handler)

images_url = ['https://storage.googleapis.com/apache-beam-samples/image_captioning/Paris-sunset.jpeg',
              'https://storage.googleapis.com/apache-beam-samples/image_captioning/Wedges.jpeg',
              'https://storage.googleapis.com/apache-beam-samples/image_captioning/Hamsters.jpeg']

# Number of captions generated per image.
NUM_CAPTIONS_PER_IMAGE = 10

# Number of top captions to display.
NUM_TOP_CAPTIONS_TO_DISPLAY = 3

with beam.Pipeline() as pipeline:

  read_images = (
            pipeline
            | "ReadUrl" >> beam.Create(images_url)
            | "ReadImages" >> beam.ParDo(ReadImagesFromUrl()))

  blip_caption_generation = (
            read_images
            | "PreprocessBlipInput" >> beam.ParDo(PreprocessBLIPInput(NUM_CAPTIONS_PER_IMAGE))
            | "GenerateCaptions" >> RunInference(BLIP_keyed_model_handler)
            | "PostprocessCaptions" >> beam.ParDo(PostprocessBLIPOutput()))

  clip_captions_ranking = (
            ({'image' : read_images, 'captions': blip_caption_generation})
            | "CreateImageCaptionPair" >> beam.CoGroupByKey()
            | "PreprocessClipInput" >> beam.ParDo(
                PreprocessCLIPInput(
                    clip_feature_extractor_config_path,
                    clip_tokenizer_vocab_config_path,
                    clip_merges_config_path))
            | "GetRankingLogits" >> RunInference(CLIP_keyed_model_handler)
            | "RankClipOutput" >> beam.ParDo(RankCLIPOutput())
            )

  clip_captions_ranking | "FormatCaptions" >> beam.ParDo(FormatCaptions(NUM_TOP_CAPTIONS_TO_DISPLAY))

Image: Paris-sunset
    Top 3 captions ranked by CLIP:
        1: the setting sun is reflected in an orange setting sky over paris. (Caption probability: 0.28)
        2: the sun rising above the eiffel tower over paris. (Caption probability: 0.23)
        3: the sun setting over the eiffel tower and rooftops. (Caption probability: 0.15)


Image: Wedges
    Top 3 captions ranked by CLIP:
        1: sweet potato fries with ketchup served in bowl. (Caption probability: 0.73)
        2: this is a plate of sweet potato fries with ketchup. (Caption probability: 0.16)
        3: sweet potato fries and a dipping sauce are on the tray. (Caption probability: 0.06)


Image: Hamsters
    Top 3 captions ranked by CLIP:
        1: person holding two small animals in their hands. (Caption probability: 0.62)
        2: a person's hand holding a small hamster in front of them. (Caption probability: 0.20)
        3: a person holding a small animal in their hands. (Caption probability: 0.09)