Ai-Instagram-Caption-Generator/app.py at main · rachit23tech/Ai-Instagram-Caption-Generator · GitHub

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
import streamlit as st
from PIL import Image
from database import save_caption

from utils import (
    generate_image_description,
    generate_caption_package,
    rank_captions,
    generate_hashtags
)

st.set_page_config(page_title="AI Instagram Caption Generator", layout="centered")

st.title("AI Instagram Caption Generator")

# Language selector
language = st.selectbox(
    "Caption Language 🌍",
    ["English", "Hindi", "Punjabi", "Spanish", "French"]
)

# Caption style selector
style = st.selectbox(
    "Caption Style 🎨",
    ["Instagram", "Funny", "Savage", "Aesthetic", "Motivational", "Poetic"]
)

uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])


if uploaded_file is not None:

    image = Image.open(uploaded_file).convert("RGB")

    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Generate captions
    if st.button("Generate Captions"):

        with st.spinner("Analyzing image..."):
            description = generate_image_description(image)
            st.session_state.description = description

        with st.spinner("Generating captions..."):

            package = generate_caption_package(description, style)

            ranked = rank_captions(description, package)

            if ranked:
                st.session_state.captions = ranked
                st.session_state.hashtags = generate_hashtags(ranked[0]["caption"])
            else:
                st.error("Failed to generate captions. Try again.")

    # Regenerate captions
    if st.button("Regenerate Captions") and "description" in st.session_state:

        with st.spinner("Generating new captions..."):

            package = generate_caption_package(st.session_state.description, style)

            ranked = rank_captions(st.session_state.description, package)

            if ranked:
                st.session_state.captions = ranked
                st.session_state.hashtags = generate_hashtags(ranked[0]["caption"])

    # Display captions
    if "captions" in st.session_state and st.session_state.captions:

        st.subheader("Generated Captions")

        st.write("Suggested Hashtags:")
        st.write(" ".join(st.session_state.hashtags))

        st.divider()

        for i, item in enumerate(st.session_state.captions):

            caption = item["caption"]

            if i == 0:
                st.markdown("⭐ **Best Caption**")

            st.markdown(f"**{i+1}. {caption}**")

            st.write("Emoji Version:", item["emoji_caption"])

            if language == "Punjabi":
                st.write("Punjabi:", item["punjabi"])

            st.write("Viral Score:", item["viral_score"], "/10")

            st.code(caption)

            rating = st.slider(
                f"Rate Caption {i+1}",
                min_value=1,
                max_value=5,
                value=3,
                key=f"rating_{i}"
            )

            if st.button(f"Save Caption {i+1}", key=f"save_{i}"):

                save_caption(uploaded_file.name, caption, rating)

                st.success("Saved to database!")

            st.divider()

        captions_text = "\n".join([c["caption"] for c in st.session_state.captions])

        st.download_button(
            "Download Captions",
            captions_text,
            file_name="captions.txt"
        )