Update app.py

app.py

@@ -17,8 +17,6 @@ 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')
@@ -73,29 +71,24 @@ 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)
+option = st.selectbox("Select input method:", ["Upload Images or Zip File"])
+
+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:
+                temp_dir = tempfile.mkdtemp()
+                zip_ref.extractall(temp_dir)
+                for file in zip_ref.namelist():
+                    if file.lower().endswith(('png', 'jpg', 'jpeg')):
+                        image_files.append(os.path.join(temp_dir, 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'):
@@ -103,13 +96,7 @@ if st.button("Generate Captions", key='generate_captions'):
         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
+            captions[image_file] = caption
         except Exception as e:
             st.error(f"Error processing {image_file}: {e}")