Upload 3 files

utils/caption_utils.py

@@ -0,0 +1,26 @@
+import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from utils.image_utils import load_image
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class ImageCaptioning:
+
+    def __int__(self):
+        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+        self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+
+    def get_caption(self, image_path):
+        image = load_image(image_path)
+
+        # Preprocessing the Image
+        img = self.processor(image, return_tensors="pt").to(device)
+
+        # Generating captions
+        output = self.model.generate(**img)
+
+        # decode the output
+        caption = self.processor.batch_decode(output, skip_special_tokens=True)[0]
+
+        return caption

utils/image_utils.py

@@ -0,0 +1,21 @@
+import requests
+from PIL import Image
+import urllib.parse as parse
+import os
+
+
+# Verify url
+def check_url(string):
+    try:
+        result = parse.urlparse(string)
+        return all([result.scheme, result.netloc, result.path])
+    except:
+        return False
+
+
+# Load an image
+def load_image(image_path):
+    if check_url(image_path):
+        return Image.open(requests.get(image_path, stream=True).raw)
+    elif os.path.exists(image_path):
+        return Image.open(image_path)

utils/topic_generation.py

@@ -0,0 +1,31 @@
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class TopicGenerator:
+
+    def __init__(self):
+        # Initialize tokenizer and model upon class instantiation
+        self.tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+        self.model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large").to(device)  # assuming you have a GPU available
+
+    def generate_topics(self, user_input, num_topics=3):
+        """
+        Generate topic sentences based on the user input.
+
+        Args:
+        - user_input (str): The input text provided by the user.
+        - num_topics (int, optional): Number of topics to generate. Defaults to 3.
+
+        Returns:
+        - list: A list of generated topic sentences.
+        """
+        prompt_text = f"Generate a topic sentence based on the following input: {user_input}"
+        input_ids = self.tokenizer(prompt_text, return_tensors="pt").input_ids.to(device)
+
+        # Generate topics
+        outputs = self.model.generate(input_ids, do_sample=True, top_k=50, temperature=0.7, max_length=50, num_return_sequences=num_topics)
+
+        # Decode the outputs and return as a list of topic sentences
+        return [self.tokenizer.decode(output, skip_special_tokens=True) for output in outputs]