Update app.py

app.py

@@ -261,6 +261,8 @@ if st.button('Check for Infringement'):
     with tab1:
         with st.spinner('Processing...'):
 
+            for path in os.listdir('/home/user/app'):
+                print(path)
             if not os.path.exists('/home/user/app/embeddings'):
                 download_db()
                 print("\u2713 Downloaded Database\n\n")
@@ -300,54 +302,54 @@ if st.button('Check for Infringement'):
                     with col2:
                         st.markdown(f"*Similar Text:* {similar_text}")
 
-    if need_image == 'True':
-        with st.spinner('Processing Images...'):
-            emb_main , main_prod_imgs = get_image_embeddings(main_product)
-            similar_prod = extract_similar_products(main_product)[0]
-            emb_similar , similar_prod_imgs = get_image_embeddings(similar_prod)
+        if need_image == 'True':
+            with st.spinner('Processing Images...'):
+                emb_main , main_prod_imgs = get_image_embeddings(main_product)
+                similar_prod = extract_similar_products(main_product)[0]
+                emb_similar , similar_prod_imgs = get_image_embeddings(similar_prod)
 
-            similarity_matrix = np.zeros((5, 5))
-            for i in range(5):
-                for j in range(5):
-                    similarity_matrix[i][j] = cosine_similarity([emb_main[i]], [emb_similar[j]])[0][0]
+                similarity_matrix = np.zeros((5, 5))
+                for i in range(5):
+                    for j in range(5):
+                        similarity_matrix[i][j] = cosine_similarity([emb_main[i]], [emb_similar[j]])[0][0]
 
-            st.subheader("Image Similarity")
-            # Create an interactive heatmap
-            fig = px.imshow(similarity_matrix,
-                            labels=dict(x=f"{similar_prod} Images", y=f"{main_product} Images", color="Similarity"),
-                            x=[f"Image {i+1}" for i in range(5)],
-                            y=[f"Image {i+1}" for i in range(5)],
-                            color_continuous_scale="Viridis")
+                st.subheader("Image Similarity")
+                # Create an interactive heatmap
+                fig = px.imshow(similarity_matrix,
+                                labels=dict(x=f"{similar_prod} Images", y=f"{main_product} Images", color="Similarity"),
+                                x=[f"Image {i+1}" for i in range(5)],
+                                y=[f"Image {i+1}" for i in range(5)],
+                                color_continuous_scale="Viridis")
 
-            # Add title to the heatmap
-            fig.update_layout(title="Image Similarity Heatmap")
+                # Add title to the heatmap
+                fig.update_layout(title="Image Similarity Heatmap")
 
-            # Display the interactive heatmap
-            st.plotly_chart(fig)
+                # Display the interactive heatmap
+                st.plotly_chart(fig)
 
 
 
-            @st.experimental_fragment
-            def image_viewer():
-                # Form to handle image selection
+                @st.experimental_fragment
+                def image_viewer():
+                    # Form to handle image selection
 
-                st.subheader("Image Viewer")
+                    st.subheader("Image Viewer")
 
-                selected_row = st.selectbox('Select a row (Main Product Image)', [f'Image {i+1}' for i in range(5)])
-                selected_col = st.selectbox('Select a column (Similar Product Image)', [f'Image {i+1}' for i in range(5)])
+                    selected_row = st.selectbox('Select a row (Main Product Image)', [f'Image {i+1}' for i in range(5)])
+                    selected_col = st.selectbox('Select a column (Similar Product Image)', [f'Image {i+1}' for i in range(5)])
 
-                # Get the selected indices from session state
-                row_idx = int(selected_row.split()[1]) - 1
-                col_idx = int(selected_col.split()[1]) - 1
+                    # Get the selected indices from session state
+                    row_idx = int(selected_row.split()[1]) - 1
+                    col_idx = int(selected_col.split()[1]) - 1
 
-                col1, col2 = st.columns(2)
+                    col1, col2 = st.columns(2)
 
-                with col1:
-                    st.image(main_prod_imgs[row_idx], caption=f'Main Product Image {row_idx+1}', use_column_width=True)
-                with col2:
-                    st.image(similar_prod_imgs[col_idx], caption=f'Similar Product Image {col_idx+1}', use_column_width=True)
+                    with col1:
+                        st.image(main_prod_imgs[row_idx], caption=f'Main Product Image {row_idx+1}', use_column_width=True)
+                    with col2:
+                        st.image(similar_prod_imgs[col_idx], caption=f'Similar Product Image {col_idx+1}', use_column_width=True)
 
-            # Call the fragment
-            image_viewer()
+                # Call the fragment
+                image_viewer()