README.md

Image captioning model

How To use this model.

Adapt the code below to your needs.

import os
from PIL import Image
import torchvision.transforms as transforms
from transformers import GPT2TokenizerFast, VisionEncoderDecoderModel

class DataProcessing:
    def __init__(self):
        # GPT-2 tokenizer
        self.tokenizer = GPT2TokenizerFast.from_pretrained('distilgpt2')
        self.tokenizer.pad_token = self.tokenizer.eos_token
        # Define the transforms to be applied to the images
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

class GenerateCaptions(DataProcessing):
    NUM_BEAMS = 3
    MAX_LENGTH = 15
    EARLY_STOPPING = True
    DO_SAMPLE = True
    TOP_K = 10
    NUM_RETURN_SEQUENCES = 2 # number of captions to generate

    def __init__(self, captioning_model):
        self.captioning_model = captioning_model
        super().__init__()

    def read_img_predict(self, path):
        try:
            with Image.open(path) as img:
                if img.mode != "RGB":
                    img = img.convert('RGB')
                img_transformed = self.transform(img).unsqueeze(0)
                # tensor dimensions max_lenght X num_return_sequences, where ij == some_token_id
                model_output = self.captioning_model.generate(
                    img_transformed,
                    num_beams=self.NUM_BEAMS,
                    max_length=self.MAX_LENGTH,
                    early_stopping=self.EARLY_STOPPING,
                    do_sample=self.DO_SAMPLE,
                    top_k=self.TOP_K,
                    num_return_sequences=self.NUM_RETURN_SEQUENCES,
                )
                # g is a tensor like this one: tensor([50256,    13,   198,   198,   198,   198,   198,   198,   198, 50256,
                # 50256, 50256, 50256, 50256, 50256])
                captions = [self.tokenizer.decode(g, skip_special_tokens=True).strip() for g in model_output]

                return captions
        except FileNotFoundError:
            raise FileNotFoundError(f"File not found: {path}")

    def generate_caption(self, path):
        """
        Generate captions for a single image or a directory of images
        :param path: path to image or directory of images
        :return: captions
        """
        if os.path.isdir(path):
            self.decoded_predictions = []
            for root, dirs, files in os.walk(path):
                for file in files:
                    self.decoded_predictions.append(self.read_img_predict(os.path.join(root, file)))
            return self.decoded_predictions
        elif os.path.isfile(path):
            return self.read_img_predict(path)
        else:
            raise ValueError(f"Invalid path: {path}")





image_captioning_model = VisionEncoderDecoderModel.from_pretrained("yesidcanoc/image-captioning-swin-tiny-distilgpt2")

generate_captions = GenerateCaptions(image_captioning_model)

captions = generate_captions.generate_caption('../data/test_data/images')

print(captions)