Update app.py

app.py

@@ -2,7 +2,8 @@ import pandas as pd
 import streamlit as st
 import datasets
 import plotly.express as px
-from sentence_transformers import SentenceTransformer, util
+from sentence_transformers import SentenceTransformer
+from PIL import Image
 import os
 from pandas.api.types import (
     is_categorical_dtype,
@@ -11,6 +12,7 @@ from pandas.api.types import (
     is_object_dtype,
 )
 import subprocess
+from tempfile import NamedTemporaryFile
 
 st.set_page_config(layout="wide")
 
@@ -30,12 +32,14 @@ def load_dataset():
     dataset = datasets.load_dataset('rjadr/ditaduranuncamais', split='train', use_auth_token=token)
     dataset.add_faiss_index(column="txt_embs")
     dataset.add_faiss_index(column="img_embs")
-    dataset = dataset.remove_columns(['Post Created Time','Like and View Counts Disabled','Link','Download URL','Views'])
+    dataset = dataset.remove_columns(['Post Created Date', 'Post Created Time','Like and View Counts Disabled','Link','Download URL','Views'])
     return dataset
 
 @st.cache_data(show_spinner=False)
 def load_dataframe(_dataset):
     dataframe = _dataset.remove_columns(['txt_embs', 'img_embs']).to_pandas()
+   # dataframe['Post Created'] = dataframe['Post Created'].dt.tz_convert('UTC')
+    dataframe = dataframe[['Post Created', 'image', 'Description', 'Image Text', 'Account', 'User Name'] + [col for col in dataframe.columns if col not in ['Post Created', 'image', 'Description', 'Image Text', 'Account', 'User Name']]]  
     return dataframe
 
 @st.cache_resource(show_spinner=True)
@@ -132,7 +136,7 @@ def get_image_embs(image):
     Returns:
     img_emb (np.array): Image embeddings
     """
-    img_emb = image_model.encode(image)
+    img_emb = image_model.encode(Image.open(image))
     return img_emb
 
 @st.cache_data(show_spinner=False)
@@ -287,71 +291,88 @@ with tab2:
 
     if selected_tab == "Text to Text":
         text_to_text_input = st.text_input("Enter text")
-        text_to_text_k_top = st.slider("Number of results", 1, 60, 8)
+        text_to_text_k_top = st.slider("Number of results", 1, 500, 8)
         if st.button("Search"):
-            st.dataframe(
-                data=text_to_text(text_to_text_input, text_to_text_k_top),
-                column_config={
-                "image": st.column_config.ImageColumn(
-                    "Image", help="Instagram image"
-                ),
-                "URL": st.column_config.LinkColumn(
-                    "Link", help="Instagram link", width="small"
-                )
-                },
-                hide_index=True,
-            )   
-            
-    elif selected_tab == "Text to Image":
-        text_to_image_input = st.text_input("Enter text")
-        text_to_image_k_top = st.slider("Number of results", 1, 60, 8)
-        if st.button("Search"):
-            st.dataframe(
-                data=text_to_image(text_to_image_input, text_to_image_k_top),
-                column_config={
+            if not text_to_text_input:
+                st.warning("Please enter text")
+            else:
+                st.dataframe(
+                    data=text_to_text(text_to_text_input, text_to_text_k_top),
+                    column_config={
                     "image": st.column_config.ImageColumn(
                         "Image", help="Instagram image"
                     ),
                     "URL": st.column_config.LinkColumn(
                         "Link", help="Instagram link", width="small"
                     )
-                },
-                hide_index=True,
-            )
+                    },
+                    hide_index=True,
+                )   
+            
+    elif selected_tab == "Text to Image":
+        text_to_image_input = st.text_input("Enter text")
+        text_to_image_k_top = st.slider("Number of results", 1, 500, 8)
+        if st.button("Search"):
+            if not text_to_image_input:
+                st.warning("Please enter some text")
+            else:
+                st.dataframe(
+                    data=text_to_image(text_to_image_input, text_to_image_k_top),
+                    column_config={
+                        "image": st.column_config.ImageColumn(
+                            "Image", help="Instagram image"
+                        ),
+                        "URL": st.column_config.LinkColumn(
+                            "Link", help="Instagram link", width="small"
+                        )
+                    },
+                    hide_index=True,
+                )
 
     elif selected_tab == "Image to Image":
-        image_to_image_k_top = st.slider("Number of results", 1, 60, 8)
+        image_to_image_k_top = st.slider("Number of results", 1, 500, 8)
         image_to_image_input = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+        temp_file = NamedTemporaryFile(delete=False)
         if st.button("Search"):
-            st.dataframe(
-                data=image_to_image(image_to_image_input, image_to_image_k_top),
-                column_config={
-                    "image": st.column_config.ImageColumn(
-                        "Image", help="Instagram image"
-                    ),
-                    "URL": st.column_config.LinkColumn(
-                        "Link", help="Instagram link", width="small"
-                    )
-                },
-                hide_index=True,
-            )
+            if not image_to_image_input:
+                st.warning("Please upload an image")
+            else:
+                temp_file.write(image_to_image_input.getvalue())
+                
+                st.dataframe(
+                    data=image_to_image(temp_file, image_to_image_k_top),
+                    column_config={
+                        "image": st.column_config.ImageColumn(
+                            "Image", help="Instagram image"
+                        ),
+                        "URL": st.column_config.LinkColumn(
+                            "Link", help="Instagram link", width="small"
+                        )
+                    },
+                    hide_index=True,
+                )
 
     elif selected_tab == "Image to Text":
-        image_to_text_k_top = st.slider("Number of results", 1, 60, 8)
+        image_to_text_k_top = st.slider("Number of results", 1, 500, 8)
         image_to_text_input = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+        temp_file = NamedTemporaryFile(delete=False)
         if st.button("Search"):
-            st.dataframe(
-                data=image_to_text(image_to_text_input, image_to_text_k_top),
-                column_config={
-                    "image": st.column_config.ImageColumn(
-                        "Image", help="Instagram image"
-                    ),
-                    "URL": st.column_config.LinkColumn(
-                        "Link", help="Instagram link", width="small"
-                    )
-                },
-                hide_index=True,
-            )   
+            if not image_to_text_input:
+                st.warning("Please upload an image")
+            else:
+                temp_file.write(image_to_text_input.getvalue())
+                st.dataframe(
+                    data=image_to_text(temp_file, image_to_text_k_top),
+                    column_config={
+                        "image": st.column_config.ImageColumn(
+                            "Image", help="Instagram image"
+                        ),
+                        "URL": st.column_config.LinkColumn(
+                            "Link", help="Instagram link", width="small"
+                        )
+                    },
+                    hide_index=True,
+                )   
             
 with tab3:
     st.markdown("### Time Series Analysis")