Update: Added estimated timer and cold-start model initialization warnings for Hugging Face Space

app.py

@@ -4,10 +4,10 @@ import os
 import sys
 import tempfile
 import pandas as pd
+import time
 
 # Ensure bacsense_v2_package is importable from the repo root
 sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
-from bacsense_v2_package.inference import BacSense
 
 # Precaution dictionary
 PRECAUTIONS = {
@@ -25,19 +25,12 @@ st.set_page_config(
     layout="wide"
 )
 
-# Initialize classifier
-@st.cache_resource
+# Lazy load — BacSense (and TensorFlow) only imported when first image is uploaded
+@st.cache_resource(show_spinner="⏳ Loading BacSense model — this may take ~60s on first run...")
 def get_classifier():
+    from bacsense_v2_package.inference import BacSense
     model_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), 'bacsense_v2_package'))
-    model = BacSense(model_dir=model_dir)
-    model.warmup()
-    return model
-
-try:
-    classifier = get_classifier()
-except Exception as e:
-    st.error(f"Error loading model: {e}")
-    st.stop()
+    return BacSense(model_dir=model_dir)
 
 # Dialog function for detailed view
 def render_details(item):
@@ -94,6 +87,9 @@ uploaded_files = st.file_uploader("Upload microscopic bacterial images...", type
 if "selected_item" not in st.session_state:
     st.session_state.selected_item = None
 
+if "is_first_run" not in st.session_state:
+    st.session_state.is_first_run = True
+
 if uploaded_files:
     if st.session_state.selected_item is not None:
         render_details(st.session_state.selected_item)
@@ -106,9 +102,24 @@ if uploaded_files:
             st.write(f"Processing {len(uploaded_files)} images...")
             progress_bar = st.progress(0)
             status_text = st.empty()
+            timer_text = st.empty()
+        
+        # Calculate estimated time: 6s per image + 120s for cold start on HF
+        per_image_time = 6
+        cold_start_time = 120 if st.session_state.is_first_run else 10
+        total_seconds = int(len(uploaded_files) * per_image_time + cold_start_time)
+        start_time = time.time()
         
         for i, uploaded_file in enumerate(uploaded_files):
-            status_text.text(f"Analyzing [{i+1}/{len(uploaded_files)}]: {uploaded_file.name}...")
+            current_elapsed = time.time() - start_time
+            remaining = max(1, total_seconds - int(current_elapsed))
+            
+            if st.session_state.is_first_run and i == 0:
+                status_text.info(f"⏳ **Initializing VGG16 Model...** (This may take up to 2 mins on first run). \n\n Analyzing: {uploaded_file.name}")
+            else:
+                status_text.text(f"Analyzing [{i+1}/{len(uploaded_files)}]: {uploaded_file.name}...")
+            
+            timer_text.markdown(f"⏱️ **Estimated time remaining: ~{remaining} seconds**")
             
             image = Image.open(uploaded_file)
             fd, temp_path = tempfile.mkstemp(suffix=".png")
@@ -119,6 +130,8 @@ if uploaded_files:
                 image.save(f, format="PNG")
             
             try:
+                # Retrieve the classifier (lazy-load) 
+                classifier = get_classifier()
                 prediction = classifier.predict(temp_path)
                 confidence_pct = prediction['confidence'] * 100 if prediction['confidence'] <= 1.0 else prediction['confidence']
                 results.append({
@@ -140,6 +153,8 @@ if uploaded_files:
             
             progress_bar.progress((i + 1) / len(uploaded_files))
         
+        st.session_state.is_first_run = False
+        timer_text.empty()
         status_text.text("Batch Processing Complete!")
         
         if results: