Update app.py

app.py

@@ -1,76 +1,141 @@
 import streamlit as st
 from PIL import Image
-from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer
-import itertools
-from nltk.corpus import stopwords
-import nltk
-import easyocr
-import torch
 import numpy as np
+import nltk
 nltk.download('stopwords')
+nltk.download('punkt')
+import pandas as pd
+import random
+import easyocr
+import re
+from nltk.corpus import stopwords
+from nltk.tokenize import word_tokenize
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
 
-# load the model and tokenizer
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+# Directory path to the saved model on Google Drive
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
+# Load the feature extractor and tokenizer
+feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-reader = easyocr.Reader(['en'])
 
-# set up Streamlit app
-st.set_page_config(layout='wide', page_title='Image Hashtag Recommender')
 
-def generate_hashtags(image_file):
-    # get image and convert to RGB mode
-    image = Image.open(image_file).convert('RGB')
+def generate_captions(image):
+    image = Image.open(image).convert("RGB")
+    generated_caption = tokenizer.decode(model.generate(feature_extractor(image, return_tensors="pt").pixel_values.to("cpu"))[0])
+    sentence = generated_caption
+    text_to_remove = "<|endoftext|>"
+    generated_caption = sentence.replace(text_to_remove, "")
+    return generated_caption
 
-    # extract image features
-    inputs = processor(image, return_tensors="pt")
+# use easyocr to extract text from the image
+def image_text(image):
+    img_np = np.array(image)
+    reader = easyocr.Reader(['en'])
+    text = reader.readtext(img_np)
+    detected_text = " ".join([item[1] for item in text])
 
-    output_ids = model.generate(**inputs)
-    # out_text = processor.decode(out[0], skip_special_tokens=True)
+    # Extract individual words, convert to lowercase, and add "#" symbol
+    detected_text= ['#' + entry[1].strip().lower().replace(" ", "") for entry in text]
+    return detected_text
 
-    # decode the model output to text and extract caption words
-    output_text = processor.decode(output_ids[0], skip_special_tokens=True)
-    caption_words = [word.lower() for word in output_text.split() if not word.startswith("#")]
+# Load NLTK stopwords for filtering
+stop_words = set(stopwords.words('english'))
 
-    # remove stop words from caption words
-    stop_words = set(stopwords.words('english'))
-    caption_words = [word for word in caption_words if word not in stop_words]
+# Add hashtags to keywords, which have been generated from image captioing
+def add_hashtags(keywords):
+    hashtags = []
+    
+    for keyword in keywords:
+        # Generate hashtag from the keyword (you can modify this part as per your requirements)
+        hashtag = '#' + keyword.lower()
+        
+        hashtags.append(hashtag)
+    
+    return hashtags
 
-    # use easyocr to extract text from the image
-    text = reader.readtext(np.array(image))
-    detected_text = " ".join([item[1] for item in text])
+def trending_hashtags(caption):
+  # Read trending hashtags from a file separated by commas
+  with open("hashies.txt", "r") as file:
+      hashtags_string = file.read()
 
-    # combine caption words and detected text
-    all_words = caption_words + detected_text.split()
+  # Split the hashtags by commas and remove any leading/trailing spaces
+  trending_hashtags = [hashtag.strip() for hashtag in hashtags_string.split(',')]
 
-    # generate combinations of words for hashtags
-    hashtags = []
-    for n in range(1, 4):
-        word_combinations = list(itertools.combinations(all_words, n))
-        for combination in word_combinations:
-            hashtag = "#" + "".join(combination)
-            hashtags.append(hashtag)
-
-    # return top 10 hashtags by frequency
-    top_hashtags = [tag for tag in sorted(set(hashtags), key=hashtags.count, reverse=True) if tag != "#"]
-    return [top_hashtags[:10], output_text]
-
-st.title("Image Caption and HashTag Recommender")
-
-image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
-# if the user has submitted an image, generate hashtags
-if image_file is not None:
-    try:
-        hashtags = generate_hashtags(image_file)
-        if len(hashtags) > 0:
-            st.write(f"Caption : {hashtags[1]}")
-            st.write("Top 10 hashtags for this image:")
-            for tag in hashtags[0]:
-                st.write(tag)
-        else:
-            st.write("No hashtags found for this image.")
-    except Exception as e:
-        st.write(f"Error: {e}")
+  # Create a DataFrame from the hashtags
+  df = pd.DataFrame(trending_hashtags, columns=["Hashtags"])
+
+  # Function to extract keywords from a given text
+  def extract_keywords(caption):
+      tokens = word_tokenize(caption)
+      keywords = [token.lower() for token in tokens if token.lower() not in stop_words]
+      return keywords
+
+  # Extract keywords from caption and trending hashtags
+  caption_keywords = extract_keywords(caption)
+  hashtag_keywords = [extract_keywords(hashtag) for hashtag in df["Hashtags"]]
+
+  # Function to calculate cosine similarity between two strings
+  def calculate_similarity(text1, text2):
+      tfidf_vectorizer = TfidfVectorizer()
+      tfidf_matrix = tfidf_vectorizer.fit_transform([text1, text2])
+      similarity_matrix = cosine_similarity(tfidf_matrix[0], tfidf_matrix[1])
+      return similarity_matrix[0][0]
+
+  # Calculate similarity between caption and each trending hashtag
+  similarities = [calculate_similarity(' '.join(caption_keywords), ' '.join(keywords)) for keywords in hashtag_keywords]
+
+  # Sort trending hashtags based on similarity in descending order
+  sorted_hashtags = [hashtag for _, hashtag in sorted(zip(similarities, df["Hashtags"]), reverse=True)]
+
+  # Select top k relevant hashtags (e.g., top 5) without duplicates
+  selected_hashtags = list(set(sorted_hashtags[:5]))
+
+  selected_hashtag = [word.strip("'") for word in selected_hashtags]
+
+  return selected_hashtag
+
+# create the Streamlit app
+def app():
+    st.title('Image from your Side, Trending Hashtags from our Side')
+
+    st.write('Upload an image to see what we have in store.')
+
+    # create file uploader
+    uploaded_file = st.file_uploader("Got You Covered, Upload your wish!, magic on the Way! ", type=["jpg", "jpeg", "png"])
+
+    # check if file has been uploaded
+    if uploaded_file is not None:
+        # load the image
+        image = Image.open(uploaded_file).convert("RGB")
+
+        # Image Captions
+        string = generate_captions(uploaded_file)
+        tokens = word_tokenize(string)
+        keywords = [token.lower() for token in tokens if token.lower() not in stop_words]
+        hashtags = add_hashtags(keywords)
+
+        # Text Captions from image
+        extracted_text = image_text(image)
+
+        #Final Hashtags Generation
+        web_hashtags = trending_hashtags(string)
+
+        combined_hashtags = hashtags + extracted_text + web_hashtags
+
+        # Shuffle the list randomly
+        random.shuffle(combined_hashtags)
+
+        combined_hashtags = list(set(item for item in combined_hashtags[:15] if not re.search(r'\d$', item)))
+
+
+        # display the image
+        st.image(image, caption='The Uploaded File')
+        st.write("First is first captions for your Photo : ", string)
+        st.write("Magical hashies have arrived : ", combined_hashtags)
 
+# run the app
+if __name__ == '__main__':
+    app()