| ''' |
| This script contains the CaptionGenerator class, |
| which is used to generate Instagram captions for images |
| using BLIP and Gemini models. |
| ''' |
| import google.generativeai as genai |
| import streamlit as st |
| import app_utils as utils |
| from transformers import AutoProcessor, Blip2ForConditionalGeneration |
|
|
| genai.configure(api_key=utils.get_gemini_api_key()) |
|
|
| class CaptionGenerator: |
| """ |
| Class for generating Instagram captions for images using BLIP and Gemini models. |
| The model from Hugging Face is used to generate the initial caption for the image, |
| which is then used as a prompt for the Gemini model to generate five distinct and |
| engaging captions for the image. |
| |
| Attributes: |
| - google_api_key (str): Google API key for accessing the Generative AI API. |
| - gemini_model (GenerativeModel): Gemini model for generating Instagram captions. |
| - blip_processor (AutoProcessor): BLIP model processor for image captioning. |
| - blip_model (BlipForConditionalGeneration): BLIP model for image captioning. |
| |
| Methods: |
| - process_image(image_data): Resize and prepare the image for caption generation. |
| - predict(image_data): Generate five Instagram captions for the provided image. |
| """ |
| def __init__(self): |
| self.gemini = genai.GenerativeModel('gemini-pro') |
| self.processor = None |
| self.model = None |
|
|
| def image_2_text(self, image_data, processor, model): |
| """ |
| Generate a caption for the provided image using the BLIP-2 model. |
| :param image_data: PIL.Image - The image for which the caption is to be generated. |
| :return: description - The description generated for the image. |
| """ |
| try: |
| inputs = processor(images=image_data, return_tensors="pt") |
| generated_ids = model.generate(**inputs, max_length=100) |
| description = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() |
| return description |
|
|
| except Exception as e: |
| st.error(f"Error occurred during image captioning: {e}") |
|
|
|
|
| def text_2_caption(self, image_description): |
| """ |
| Generate five Instagram captions for the provided image. |
| The image is first processed before generating captions. |
| |
| :param image_description: str - The description of the image for which captions are to be generated. |
| :return: str - Five Instagram captions formatted as specified. |
| """ |
| prompt = ( |
| f"This caption was generated with a deep learning model." |
| f"Your task is to enhance the caption to make it more engaging:" |
| f"Given this provided photo description, generate five distinct " |
| f"fun and engaging Instagram captions. Each caption must include " |
| f"at least one emoji and one hashtag. The captions should be " |
| f"formatted with a preceding 'Caption #', followed by the " |
| f"caption text. Ensure each caption is separated by a blank " |
| f"line for readability." |
| f"Original Caption: {image_description}" |
| f"Please format your response as follows: \n" |
| f"**Caption 1**: [caption text]\n" |
| f"**Caption 2**: [caption text]\n" |
| f"**Caption 3**: [caption text]\n" |
| f"**Caption 4**: [caption text]\n" |
| f"**Caption 5**: [caption text]\n" |
| ) |
|
|
| try: |
| response = self.gemini.generate_content(prompt) |
| caption = response.parts[0].text |
| caption_list = response.parts[0].text.split("\n") |
| return caption, caption_list, image_description |
|
|
| except Exception as e: |
| st.error(f"Unable to connect to Gemini API: {e}") |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| def caption_2_hashtag(self, caption): |
| """ |
| Generate additional hashtags based on the content of the caption. |
| |
| :param caption: str - The caption for which hashtags are to be generated. |
| :return: str - Additional hashtags based on the content of the caption. |
| """ |
| |
| prompt = (f"Given the provided caption, generate relevant hashtags to increase engagement," |
| f"and are related to the caption content. Original Image Description: {caption}," |
| f"Please format your response as follows:\n" |
| f'[hashtags separated by commas]' |
| f" \n") |
|
|
| try: |
| response = self.gemini.generate_content(prompt) |
| hashtags = response.parts[0].text |
| return hashtags |
|
|
| except Exception as e: |
| st.error(f"Error occurred with Gemini API: {e}") |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
