Update app.py

changes made to reflect app2

app.py

@@ -1,98 +1,117 @@
 import streamlit as st
-from PIL import Image, ImageOps
-import numpy as np
-import pickle
+from PIL import Image
 from torchvision import transforms
+from transformers import AutoProcessor, BlipForConditionalGeneration
+import google.generativeai as genai
+
+# Define model
+
+# Initialize the image to caption model - BLIP
+blip_processor = AutoProcessor.\
+    from_pretrained("Salesforce/blip-image-captioning-base")
+blip_model = BlipForConditionalGeneration.\
+    from_pretrained("Salesforce/blip-image-captioning-base")
+
+# Initialize the caption to instagram post model - gemini
+GOOGLE_API_KEY = 'AIzaSyDMeFOnOflzYJ-cjedJ8ky9AE-yThcNXFk'
+genai.configure(api_key=GOOGLE_API_KEY)
+gemini_model = genai.GenerativeModel('gemini-pro')
+
+
+def import_and_predict(image_data):
+    """Creates five possible Instagram captions for an image.
+    Args:
+    - image_data: Image data in RGB format.
+    Returns:
+    - response.text: The five Instagram captions in text format.
+    """
+    transform = transforms.Compose([
+        transforms.Resize((1080, 1080))
+    ])
+    uploaded_image = transform(image_data)
 
-
-# use $ streamlit run app.py to run app!
-
-# hide deprication warnings
-import warnings
-warnings.filterwarnings("ignore")
-
-# BLIP Model
-from transformers import BlipForConditionalGeneration, BlipProcessor
-blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-
-# Caption Model
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-caption_tokenizer = AutoTokenizer.from_pretrained("prasanthsagirala/text-to-social-media-captions")
-caption_model = AutoModelForSeq2SeqLM.from_pretrained("prasanthsagirala/text-to-social-media-captions")
-
-# with open("models/blip_model.pkl", "rb") as f:
-#     blip_model = pickle.load(f)
-
-# with open("models/caption_tokenizer.pkl", "rb") as f:
-#     caption_tokenizer = pickle.load(f)
-
-# with open("models/caption_model.pkl", "rb") as f:
-#     caption_model = pickle.load(f)
-
+    # BLIP Description Generation
+    inputs = blip_processor(images=uploaded_image, return_tensors="pt")
+    generated_ids = blip_model.generate(**inputs,
+                                        max_new_tokens=100,
+                                        max_length=100)
+    caption = blip_processor.\
+        batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+
+    # Gemini Caption Generation
+    prompt = f"""Given the provided photo caption, generate five distinct \
+                and engaging Instagram captions. Each caption must include \
+                at least one emoji and one hashtag. The captions should be \
+                formatted with a preceding "Caption #", followed by the \
+                caption text. Ensure each caption is seperated by a blank \
+                line for readability. \
+
+                Original Caption: {caption}
+
+                Please format your response as follows:
+                **Caption 1**: [caption text] \n
+                **Caption 2**: [caption text] \n
+                **Caption 3**: [caption text] \n
+                **Caption 4**: [caption text] \n
+                **Caption 5**: [caption text] \n
+                """
+    response = gemini_model.generate_content(prompt)
+    return response.text
+
+
+# Define streamlit configurations
 
 # Set pre-defined page configurations
-st.set_page_config(
-    page_title="Instamuse",  # Title
-    page_icon=":camera:",  # log-icon
-    initial_sidebar_state='auto'  # page loading state
-)
+st.set_page_config(page_title="Instamuse",
+                   page_icon=":camera:",
+                   initial_sidebar_state='auto')
 
-# Sidebard (left side of page)
+# Sidebar (left side of the page)
 with st.sidebar:
     st.image('insta.png')
     st.title("InstaMuse")
-    st.subheader(
-        "Welcome to InstaMuse, the ultimate tool for turning your snapshots into social media sensations!")
-    st.write(
-        "Start turning heads with your posts. Use InstaMuse now and watch your likes soar! ")
+    st.subheader("Welcome to InstaMuse, the ultimate tool for turning your \
+                 snapshots into social media sensations!")
+    st.write("Start turning heads with your posts. Use InstaMuse now and \
+             watch your likes soar! ")
 
 # Main page text
 st.write("""
          # InstaMuse  🌟📸
 
-        Struggling to find the perfect words to match your pictures? Let InstaMuse do the heavy \
-        lifting! With just a simple upload, our app uses cutting-edge technology to analyze your \
-        image and generate a witty, engaging, or inspiring caption that captures the essence of \
+        Struggling to find the perfect words to match your pictures? Let \
+        InstaMuse do the heavy lifting! With just a simple upload, our app \
+        uses cutting-edge technology to analyze your image and generate a \
+        witty, engaging, or inspiring caption that captures the essence of \
         your moment. \n \
 
-        Whether you’re a selfie savant, a nature explorer, or a foodie fanatic, InstaMuse is here to \
-        amplify your Instagram presence. Jazz up your feed with tailored captions that resonate with your \
-        followers and attract new eyes to your profile. It’s quick, easy, and fun!
+        Whether you’re a selfie savant, a nature explorer, or a foodie fanatic\
+        , InstaMuse is here to amplify your Instagram presence. Jazz up your \
+        feed with tailored captions that resonate with your followers and \
+        attract new eyes to your profile. It’s quick, easy, and fun!
 
         **Drag your photo here and spark some caption magic!** ✨
          """
          )
 
-# For Bhumika!!!
-# Modeling part
+# Upload image file and process image
 file = st.file_uploader("", type=["jpg", "png"])
 
-
-def import_and_predict(image_data):  # Will also need to import model I think?
-    transform = transforms.Compose([
-        transforms.Resize((1080, 1080))
-    ])
-    image = transform(image_data)
-
-    # BLIP Description Generation
-    inputs = blip_processor(images=image, return_tensors="pt")
-    generated_ids = blip_model.generate(**inputs, max_new_tokens=50)
-    generated_text = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-
-    # Caption Generation
-    inputs = ["Instagram captionize:" + generated_text]
-    inputs = caption_tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
-    output = caption_model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=64)
-    decoded_output = caption_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
-    return decoded_output
-
-
-if file is None:
-    st.text("Upload your photo now and let the caption fun begin!")
+if file is not None:
+    # Create two columns for the image and the captions
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        # Image column, left screen.
+        st.markdown("#### Photo:")
+        image = Image.open(file).convert('RGB')
+        image.thumbnail((400, 400))
+        st.image(image, caption='Uploaded Image')
+
+    with col2:
+        # Caption column, right screen.
+        predictions = import_and_predict(image)
+        st.markdown("#### Captions:")
+        st.write(predictions)
 else:
-    image = Image.open(file).convert('RGB')
-    st.image(image, use_column_width=True)
-    predictions = import_and_predict(image)
-    st.markdown("## Captions:")
-    st.info(predictions)
+    st.text("Please upload an image to generate captions.")