Update app.py

app.py

@@ -78,26 +78,33 @@ st.title("Image Gallery with Captioning and Search")
 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:", ["Upload Images or Zip File"])
+# Options for input strategy
+input_option = st.selectbox("Select input method:", ["Folder Path", "Upload Images", "Upload ZIP"])
 
-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])
+
+if input_option == "Folder Path":
+    folder_path = st.text_input("Enter the folder path containing images:")
+    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'))]
+
+elif input_option == "Upload Images":
+    uploaded_files = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+    if uploaded_files:
+        for uploaded_file in uploaded_files:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as temp_file:
                 temp_file.write(uploaded_file.read())
                 image_files.append(temp_file.name)
 
+elif input_option == "Upload ZIP":
+    uploaded_zip = st.file_uploader("Upload a ZIP file containing images", type=["zip"])
+    if uploaded_zip:
+        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+            temp_file.write(uploaded_zip.read())
+            with zipfile.ZipFile(temp_file.name, 'r') as zip_ref:
+                zip_ref.extractall("/tmp/images")
+                image_files = [os.path.join("/tmp/images", f) for f in zip_ref.namelist() if f.lower().endswith(('png', 'jpg', 'jpeg'))]
+
 captions = {}
 if st.button("Generate Captions", key='generate_captions'):
     for image_file in image_files: