test / app.py
witchEverly's picture
Update app.py
db3cf05 verified
import streamlit as st
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)
# 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",
page_icon=":camera:",
initial_sidebar_state='auto')
# 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! ")
# 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 \
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!
**Drag your photo here and spark some caption magic!** ✨
"""
)
# Upload image file and process image
file = st.file_uploader("", type=["jpg", "png"])
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:
st.text("Please upload an image to generate captions.")