Update app.py

app.py

@@ -1,162 +1,171 @@
-import streamlit as st
-import os
-import zipfile
-import tempfile
-import base64
-from PIL import Image
-from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
-import pandas as pd
-from nltk.corpus import wordnet
-import spacy
-import io
-
-# Download NLTK WordNet data
-import nltk
-nltk.download('wordnet')
-nltk.download('omw-1.4')
-
-# Load spaCy model
-nlp = spacy.load("en_core_web_sm")
-
-# Load the pre-trained model for image captioning
-model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-85k-09"
-model = VisionEncoderDecoderModel.from_pretrained(model_name)
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-def generate_caption(image):
-    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-    output_ids = model.generate(pixel_values)
-    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return caption
-
-def get_synonyms(word):
-    synonyms = set()
-    for syn in wordnet.synsets(word):
-        for lemma in syn.lemmas():
-            synonyms.add(lemma.name())
-    return synonyms
-
-def preprocess_query(query):
-    doc = nlp(query)
-    tokens = set()
-    for token in doc:
-        tokens.add(token.text)
-        tokens.add(token.lemma_)
-        tokens.update(get_synonyms(token.text))
-    return tokens
-
-def search_captions(query, captions):
-    query_tokens = preprocess_query(query)
-    
-    results = []
-    for path, caption in captions.items():
-        caption_tokens = preprocess_query(caption)
-        if query_tokens & caption_tokens:
-            results.append((path, caption))
-    
-    return results
-
-st.title("Image Gallery with Captioning and Search")
-
-# Sidebar for search functionality
-with st.sidebar:
-    query = st.text_input("Search images by caption:")
-
-# Right side for folder path input and displaying images
-option = st.selectbox("Select input method:", ["Folder Path", "Upload Images"])
-
-if option == "Folder Path":
-    folder_path = st.text_input("Enter the folder path containing images:")
-    image_files = []
-    if folder_path and os.path.isdir(folder_path):
-        image_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.lower().endswith(('png', 'jpg', 'jpeg'))]
-else:
-    uploaded_files = st.file_uploader("Upload images or a zip file containing images:", type=['png', 'jpg', 'jpeg', 'zip'], accept_multiple_files=True)
-    image_files = []
-    if uploaded_files:
-        for uploaded_file in uploaded_files:
-            if uploaded_file.name.endswith('.zip'):
-                with zipfile.ZipFile(uploaded_file, 'r') as zip_ref:
-                    zip_ref.extractall("uploaded_images")
-                    for file in zip_ref.namelist():
-                        if file.lower().endswith(('png', 'jpg', 'jpeg')):
-                            image_files.append(os.path.join("uploaded_images", file))
-            else:
-                if uploaded_file.name.lower().endswith(('png', 'jpg', 'jpeg')):
-                    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1])
-                    temp_file.write(uploaded_file.read())
-                    image_files.append(temp_file.name)
-
-captions = {}
-if st.button("Generate Captions", key='generate_captions'):
-    for image_file in image_files:
-        try:
-            image = Image.open(image_file)
-            caption = generate_caption(image)
-            if option == "Folder Path":
-                captions[os.path.join(folder_path, os.path.basename(image_file))] = caption
-            else:
-                if image_file.startswith("uploaded_images"):
-                    captions[image_file.replace("uploaded_images/", "")] = caption
-                else:
-                    captions[os.path.basename(image_file)] = caption
-        except Exception as e:
-            st.error(f"Error processing {image_file}: {e}")
-
-    # Display images in a grid
-    st.subheader("Images and Captions:")
-    cols = st.columns(4)
-    idx = 0
-    for image_path, caption in captions.items():
-        col = cols[idx % 4]
-        with col:
-            try:
-                with open(image_path, "rb") as img_file:
-                    img_bytes = img_file.read()
-                encoded_image = base64.b64encode(img_bytes).decode()
-                st.markdown(
-                    f"""
-                    <div style='text-align: center;'>
-                        <img src='data:image/jpeg;base64,{encoded_image}' width='100%'>
-                        <p>{caption}</p>
-                        <p style='font-size: small; font-style: italic;'>{image_path}</p>
-                    </div>
-                    """, unsafe_allow_html=True)
-            except Exception as e:
-                st.error(f"Error displaying {image_path}: {e}")
-        idx += 1
-
-    if query:
-        results = search_captions(query, captions)
-        st.write("Search Results:")
-        cols = st.columns(4)
-        idx = 0
-        for image_path, caption in results:
-            col = cols[idx % 4]
-            with col:
-                try:
-                    with open(image_path, "rb") as img_file:
-                        img_bytes = img_file.read()
-                    encoded_image = base64.b64encode(img_bytes).decode()
-                    st.markdown(
-                        f"""
-                        <div style='text-align: center;'>
-                            <img src='data:image/jpeg;base64,{encoded_image}' width='100%'>
-                            <p>{caption}</p>
-                            <p style='font-size: small; font-style: italic;'>{image_path}</p>
-                        </div>
-                        """, unsafe_allow_html=True)
-                except Exception as e:
-                    st.error(f"Error displaying search result {image_path}: {e}")
-            idx += 1
-
-    # Save captions to Excel and provide a download button
-    df = pd.DataFrame(list(captions.items()), columns=['Image', 'Caption'])
-    excel_file = io.BytesIO()
-    df.to_excel(excel_file, index=False)
-    excel_file.seek(0)
-    st.download_button(label="Download captions as Excel",
-                       data=excel_file,
-                       file_name="captions.xlsx",
-                       mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
+import streamlit as st
+import os
+import zipfile
+import tempfile
+import base64
+from PIL import Image
+from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
+import pandas as pd
+from nltk.corpus import wordnet
+import spacy
+import io
+from spacy.cli import download
+
+# Download the model if it is not already present
+download("en_core_web_sm")
+
+# Load the model
+nlp = spacy.load("en_core_web_sm")
+
+# Your existing code here
+
+# Download NLTK WordNet data
+import nltk
+nltk.download('wordnet')
+nltk.download('omw-1.4')
+
+# Load spaCy model
+nlp = spacy.load("en_core_web_sm")
+
+# Load the pre-trained model for image captioning
+model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-85k-09"
+model = VisionEncoderDecoderModel.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+def generate_caption(image):
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+    output_ids = model.generate(pixel_values)
+    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return caption
+
+def get_synonyms(word):
+    synonyms = set()
+    for syn in wordnet.synsets(word):
+        for lemma in syn.lemmas():
+            synonyms.add(lemma.name())
+    return synonyms
+
+def preprocess_query(query):
+    doc = nlp(query)
+    tokens = set()
+    for token in doc:
+        tokens.add(token.text)
+        tokens.add(token.lemma_)
+        tokens.update(get_synonyms(token.text))
+    return tokens
+
+def search_captions(query, captions):
+    query_tokens = preprocess_query(query)
+    
+    results = []
+    for path, caption in captions.items():
+        caption_tokens = preprocess_query(caption)
+        if query_tokens & caption_tokens:
+            results.append((path, caption))
+    
+    return results
+
+st.title("Image Gallery with Captioning and Search")
+
+# Sidebar for search functionality
+with st.sidebar:
+    query = st.text_input("Search images by caption:")
+
+# Right side for folder path input and displaying images
+option = st.selectbox("Select input method:", ["Folder Path", "Upload Images"])
+
+if option == "Folder Path":
+    folder_path = st.text_input("Enter the folder path containing images:")
+    image_files = []
+    if folder_path and os.path.isdir(folder_path):
+        image_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.lower().endswith(('png', 'jpg', 'jpeg'))]
+else:
+    uploaded_files = st.file_uploader("Upload images or a zip file containing images:", type=['png', 'jpg', 'jpeg', 'zip'], accept_multiple_files=True)
+    image_files = []
+    if uploaded_files:
+        for uploaded_file in uploaded_files:
+            if uploaded_file.name.endswith('.zip'):
+                with zipfile.ZipFile(uploaded_file, 'r') as zip_ref:
+                    zip_ref.extractall("uploaded_images")
+                    for file in zip_ref.namelist():
+                        if file.lower().endswith(('png', 'jpg', 'jpeg')):
+                            image_files.append(os.path.join("uploaded_images", file))
+            else:
+                if uploaded_file.name.lower().endswith(('png', 'jpg', 'jpeg')):
+                    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1])
+                    temp_file.write(uploaded_file.read())
+                    image_files.append(temp_file.name)
+
+captions = {}
+if st.button("Generate Captions", key='generate_captions'):
+    for image_file in image_files:
+        try:
+            image = Image.open(image_file)
+            caption = generate_caption(image)
+            if option == "Folder Path":
+                captions[os.path.join(folder_path, os.path.basename(image_file))] = caption
+            else:
+                if image_file.startswith("uploaded_images"):
+                    captions[image_file.replace("uploaded_images/", "")] = caption
+                else:
+                    captions[os.path.basename(image_file)] = caption
+        except Exception as e:
+            st.error(f"Error processing {image_file}: {e}")
+
+    # Display images in a grid
+    st.subheader("Images and Captions:")
+    cols = st.columns(4)
+    idx = 0
+    for image_path, caption in captions.items():
+        col = cols[idx % 4]
+        with col:
+            try:
+                with open(image_path, "rb") as img_file:
+                    img_bytes = img_file.read()
+                encoded_image = base64.b64encode(img_bytes).decode()
+                st.markdown(
+                    f"""
+                    <div style='text-align: center;'>
+                        <img src='data:image/jpeg;base64,{encoded_image}' width='100%'>
+                        <p>{caption}</p>
+                        <p style='font-size: small; font-style: italic;'>{image_path}</p>
+                    </div>
+                    """, unsafe_allow_html=True)
+            except Exception as e:
+                st.error(f"Error displaying {image_path}: {e}")
+        idx += 1
+
+    if query:
+        results = search_captions(query, captions)
+        st.write("Search Results:")
+        cols = st.columns(4)
+        idx = 0
+        for image_path, caption in results:
+            col = cols[idx % 4]
+            with col:
+                try:
+                    with open(image_path, "rb") as img_file:
+                        img_bytes = img_file.read()
+                    encoded_image = base64.b64encode(img_bytes).decode()
+                    st.markdown(
+                        f"""
+                        <div style='text-align: center;'>
+                            <img src='data:image/jpeg;base64,{encoded_image}' width='100%'>
+                            <p>{caption}</p>
+                            <p style='font-size: small; font-style: italic;'>{image_path}</p>
+                        </div>
+                        """, unsafe_allow_html=True)
+                except Exception as e:
+                    st.error(f"Error displaying search result {image_path}: {e}")
+            idx += 1
+
+    # Save captions to Excel and provide a download button
+    df = pd.DataFrame(list(captions.items()), columns=['Image', 'Caption'])
+    excel_file = io.BytesIO()
+    df.to_excel(excel_file, index=False)
+    excel_file.seek(0)
+    st.download_button(label="Download captions as Excel",
+                       data=excel_file,
+                       file_name="captions.xlsx",
+                       mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")