spliting major tom europe in smaller countries

app.py

@@ -1,62 +1,59 @@
 import streamlit as st
-import logging
-import os
-import time
-import psutil
-from helper import (
+from helper3 import (
     load_dataset, search, get_file_paths,
     get_cordinates, get_images_from_s3_to_display,
     get_images_with_bounding_boxes_from_s3, load_dataset_with_limit
 )
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+import os
+import time
+import psutil
+from memory_profiler import memory_usage
 
 # Load environment variables
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
 
 # Predefined list of datasets
-datasets = ["WayveScenes", "MajorTom-Europe"]
+datasets = ["WayveScenes", "MajorTom-Germany"]
 description = {
-    "StopSign_test": "A test dataset for me",
-    "WayveScenes": "A large-scale dataset featuring diverse urban driving scenes.",
-    "MajorTom-Europe": "A geospatial dataset containing satellite imagery from across Europe."
+    "WayveScenes": "A large-scale dataset featuring diverse urban driving scenes, captured from autonomous vehicles to advance AI perception and navigation in complex environments.",
+    "MajorTom-Germany": "A geospatial dataset containing satellite imagery from across Germany, designed for tasks like land-use classification, environmental monitoring, and earth observation analytics."
 }
 selection = {
     'WayveScenes': [1, 8],
-    "MajorTom-Europe": [1, 18]
+    "MajorTom-Germany": [1, 1]
 }
-
+folder_path_dict = {
+        "WayveScenes" : 'WayveScenes/',
+        "MajorTom-Germany": "MajorTom-Europe/"
+    }
 # AWS S3 bucket name
 bucket_name = "datasets-quasara-io"
 
-# Function to log CPU and memory usage
-def log_resource_usage(stage):
-    cpu_usage = psutil.cpu_percent(interval=1)
-    memory_info = psutil.virtual_memory()
-    logging.info(f"{stage} - CPU Usage: {cpu_usage}%, Memory Usage: {memory_info.percent}%")
+# Function to display CPU and memory usage
+def display_usage():
+    process = psutil.Process(os.getpid())
+    st.write(f"CPU usage: {process.cpu_percent()}%")
+    st.write(f"Memory usage: {process.memory_info().rss / (1024 ** 2)} MB")
 
 # Streamlit App
 def main():
     # Initialize session state variables if not already initialized
     if 'search_in_small_objects' not in st.session_state:
         st.session_state.search_in_small_objects = False
+
     if 'dataset_number' not in st.session_state:
         st.session_state.dataset_number = 1
+
     if 'df' not in st.session_state:
         st.session_state.df = None
 
     st.title("Semantic Search and Image Display")
-    log_resource_usage("Initialization")
 
     # Select dataset from dropdown
     dataset_name = st.selectbox("Select Dataset", datasets)
 
-    if dataset_name == 'StopSign_test':
-        folder_path = ""
-    else:
-        folder_path = f'{dataset_name}/'
+    folder_path = folder_path_dict[dataset_name]
 
     st.caption(description[dataset_name])
 
@@ -64,13 +61,15 @@ def main():
         st.session_state.search_in_small_objects = True
         st.text("Small Object Search Enabled")
         st.session_state.dataset_number = st.selectbox("Select Subset of Data", list(range(1, selection[dataset_name][1] + 1)))
+        st.text(f"You have selected Split Dataset {st.session_state.dataset_number}")
     else:
         st.session_state.search_in_small_objects = False
         st.text("Small Object Search Disabled")
         st.session_state.dataset_number = st.selectbox("Select Subset of Data", list(range(1, selection[dataset_name][0] + 1)))
+        st.text(f"You have selected Main Dataset {st.session_state.dataset_number}")
 
-    dataset_limit = st.slider("Size of Dataset to be searched from", min_value=1000, max_value=20000, value=10000)
-    st.text(f'The smaller the dataset, the faster the search will work.')
+    dataset_limit = st.slider("Size of Dataset to be searched from", min_value=1000, max_value=30000, value=10000)
+    st.text(f'The smaller the dataset the faster the search will work.')
 
     # Load dataset with limit only if not already loaded
     if st.button("Load Dataset"):
@@ -78,25 +77,32 @@ def main():
             loading_dataset_text = st.empty()
             loading_dataset_text.text("Loading Dataset...")
             loading_dataset_bar = st.progress(0)
-
+            
+            # Memory profiling
+            mem_usage = memory_usage((load_dataset_with_limit, (dataset_name, st.session_state.dataset_number, st.session_state.search_in_small_objects), {"limit": dataset_limit}))
+            st.write(f"Memory used for loading the dataset: {mem_usage[-1]:.2f} MB")
+            
             # Simulate dataset loading progress
             for i in range(0, 100, 25):
-                time.sleep(0.2)
+                time.sleep(0.2)  # Simulate work being done
                 loading_dataset_bar.progress(i + 25)
 
-            log_resource_usage("Before Loading Dataset")
+            # Load dataset and monitor CPU and memory
             df, total_rows = load_dataset_with_limit(dataset_name, st.session_state.dataset_number, st.session_state.search_in_small_objects, limit=dataset_limit)
+            
+            # Store loaded dataset in session state
             st.session_state.df = df
-
             loading_dataset_bar.progress(100)
             loading_dataset_text.text("Dataset loaded successfully!")
             st.success(f"Dataset loaded successfully with {len(df)} rows.")
-            log_resource_usage("After Loading Dataset")
-
+            
+            # Display CPU and memory usage
+            display_usage()
+            
         except Exception as e:
-            logging.error(f"Failed to load dataset: {e}")
             st.error(f"Failed to load dataset: {e}")
-
+    
+    
     # Input search query
     query = st.text_input("Enter your search query")
 
@@ -110,23 +116,25 @@ def main():
             st.warning("Please enter a search query.")
         else:
             try:
+                # Progress bar for search
                 search_loading_text = st.empty()
                 search_loading_text.text("Searching...")
                 search_progress_bar = st.progress(0)
 
-                log_resource_usage("Before Search")
+                # Perform search on the loaded dataset from session state
                 df = st.session_state.df
                 if st.session_state.search_in_small_objects:
                     results = search(query, df, limit)
                     top_k_paths = get_file_paths(df, results)
                     top_k_cordinates = get_cordinates(df, results)
                 else:
+                    # Normal Search
                     results = search(query, df, limit)
                     top_k_paths = get_file_paths(df, results)
 
+                # Complete the search progress
                 search_progress_bar.progress(100)
                 search_loading_text.text("Search completed!")
-                log_resource_usage("After Search")
 
                 # Load Images with Bounding Boxes if applicable
                 if st.session_state.search_in_small_objects and top_k_paths and top_k_cordinates:
@@ -134,11 +142,14 @@ def main():
                 elif not st.session_state.search_in_small_objects and top_k_paths:
                     st.write(f"Displaying top {len(top_k_paths)} results for query '{query}':")
                     get_images_from_s3_to_display(bucket_name, top_k_paths, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, folder_path)
+                    
                 else:
                     st.write("No results found.")
 
+                # Display CPU and memory usage
+                display_usage()
+
             except Exception as e:
-                logging.error(f"Search failed: {e}")
                 st.error(f"Search failed: {e}")
 
 if __name__ == "__main__":