File size: 6,548 Bytes
73e907d 0aff634 73e907d 0aff634 73e907d 0aff634 73e907d 0aff634 73e907d e5b1de1 73e907d 0aff634 73e907d e5b1de1 0aff634 4242100 73e907d 0aff634 73e907d e0f7730 0aff634 73e907d 0aff634 73e907d 0aff634 e5b1de1 0aff634 4242100 73e907d 0aff634 173a9b9 73e907d 0aff634 e5b1de1 0aff634 73e907d b466286 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import sqlite3
import streamlit as st
from passlib.hash import bcrypt
from PIL import Image
from model import get_caption_model, generate_caption
from googletrans import Translator
import requests
# Initialize Streamlit app
st.set_page_config(page_title="Image Caption Generator", layout="wide")
# Initialize Translator
translator = Translator()
# Function to create the SQLite table if it doesn't exist
@st.cache_resource
def create_table():
with sqlite3.connect("login.db") as conn:
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL UNIQUE,
password TEXT NOT NULL,
email TEXT NOT NULL,
role TEXT NOT NULL
)
''')
# Function to handle user signup
def signup():
st.markdown("<p style='font-size: 24px; font-weight: bold; margin-bottom: 20px; text-align: center;'>Signup</p>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center;'>Please fill in the details to sign up:</p>", unsafe_allow_html=True)
new_username = st.text_input("New Username")
new_password = st.text_input("New Password", type="password")
new_email = st.text_input("Email")
if st.button("Signup"):
if not new_username or not new_password or not new_email:
st.error("All fields are required for signup.")
return
role = "user"
hashed_password = new_password
try:
with sqlite3.connect("login.db") as conn:
cursor = conn.cursor()
cursor.execute("INSERT INTO users (username, password, email, role) VALUES (?, ?, ?, ?)",
(new_username, hashed_password, new_email, role))
st.success("Signup successful! You can now login.")
except sqlite3.IntegrityError:
st.error("Username already exists. Please choose a different username.")
# Function to handle user login
def login():
st.markdown("<p style='font-size: 24px; font-weight: bold; margin-bottom: 20px; text-align: center;'>Login</p>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center;'>Please enter your login details:</p>", unsafe_allow_html=True)
username = st.text_input("Username")
password = st.text_input("Password", type="password")
if st.button("Login"):
if not username or not password:
st.error("Username and password are required for login.")
return
try:
with sqlite3.connect("login.db") as conn:
cursor = conn.cursor()
cursor.execute("SELECT * FROM users WHERE username = ?", (username,))
user = cursor.fetchone()
if user and new_password:
st.success("Login successful!")
st.write(f"You are logged in as: {user[1]}")
st.session_state.username = username
st.session_state.selected_tab = "Generate Caption"
else:
st.error("Login failed. Invalid username or password.")
except sqlite3.OperationalError as e:
st.error(f"An error occurred while trying to log in: {e}")
# Function to generate image caption and edit captions
@st.cache_resource
def generate_image_caption(img, caption_model):
return generate_caption(img, caption_model)
def display_edit_caption(selected_languages, generated_caption):
st.markdown("<p style='font-size: 24px; font-weight: bold; margin-bottom: 20px;'>Edit Caption:</p>", unsafe_allow_html=True)
edited_caption = st.text_area("Edit the caption", value=generated_caption, key="edited_caption")
if edited_caption:
st.markdown("<p style='font-size: 24px; font-weight: bold; margin-bottom: 20px;'>Edited Caption:</p>", unsafe_allow_html=True)
st.write(edited_caption)
for lang in selected_languages:
if lang != "en":
translated_caption = translator.translate(edited_caption, src="en", dest=lang)
st.markdown(f"<p style='font-size: 24px; font-weight: bold; margin-bottom: 20px;'>{lang.upper()} Translation:</p>", unsafe_allow_html=True)
st.write(translated_caption.text)
username = st.session_state.username
st.balloons()
st.success("Caption editing complete!")
# Main function to control the application flow
def main():
create_table()
tabs = ["Signup", "Login", "Generate Caption"]
selected_tab = st.sidebar.selectbox("Navigation", tabs)
if selected_tab == "Signup":
signup()
elif selected_tab == "Login":
login()
elif selected_tab == "Generate Caption":
if hasattr(st.session_state, "username"):
st.sidebar.info("Welcome to the Image Caption Generator!")
st.sidebar.warning("Be sure to upload a clear and relevant image.")
generate_caption_button = st.sidebar.button("Generate Caption")
if generate_caption_button:
st.sidebar.info("Generating caption... Please wait.")
with st.spinner("Generating caption..."):
img_url = st.sidebar.text_input("Enter Image URL:")
img_upload = st.sidebar.file_uploader("Upload Image:", type=['jpg', 'png', 'jpeg'])
if img_url or img_upload:
if img_url:
img = Image.open(requests.get(img_url, stream=True).raw)
else:
img = Image.open(img_upload)
img = img.convert('RGB')
caption_model = get_caption_model()
generated_caption = generate_image_caption(img, caption_model)
if generated_caption:
st.sidebar.success("Caption generated successfully!")
selected_languages = st.sidebar.multiselect("Select languages for translation:", ['en', 'ta', 'hi', 'zh-cn', 'es', 'fr', 'de', 'it', 'ja'])
display_edit_caption(selected_languages, generated_caption)
else:
st.sidebar.error("Caption generation failed.")
else:
generate_image_caption()
else:
st.write("Please login to access this feature.")
if __name__ == "__main__":
main() |