Update app.py

app.py

@@ -345,6 +345,17 @@ def get_model():
 
 caption_model = get_model()
 
+@st.cache_data
+def extract_important_term(caption):
+    # Remove stopwords
+    stop_words = set(stopwords.words('english'))
+    words = caption.lower().split()
+    filtered_words = [word for word in words if word not in stop_words]
+
+    # Find the longest word
+    important_term = filtered_words
+
+    return important_term
 
 
 def generate_poem(word, num_lines):
@@ -367,11 +378,14 @@ def generate_poem(word, num_lines):
 
 def predict(term_col, poem_col):
     pred_caption = generate_caption('tmp.jpg', caption_model)
+    # Extract the important term
+    important_term = extract_important_term(pred_caption)
     
     # Generate poem using poetpy
-    poem_lines = generate_poem(pred_caption, 10)
+    poem_lines = generate_poem(extract_important_term, num_lines=10)
 
     # Display the poem
+    
     poem_col.markdown('#### Generated Poem:')
     poem_col.markdown('<div class="poem-container">', unsafe_allow_html=True)
     for line in poem_lines:
@@ -379,7 +393,6 @@ def predict(term_col, poem_col):
     poem_col.markdown('</div>', unsafe_allow_html=True)
 
 
-
    
 st.markdown('<h1 style="text-align:center; font-family:Comic Sans MS; width:fit-content; font-size:3em; color:green; text-shadow: 2px 2px 4px #000000;">AUTO POEM GENERATOR</h1>', unsafe_allow_html=True)
 col1, col2 = st.columns(2)