Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,136 +1,143 @@
 import streamlit as st
 import numpy as np
+from PIL import Image
 import cv2
 import os
 import sys
-from PIL import Image
-import traceback
 
-# Configure the app
+# Set page configuration
 st.set_page_config(
-    page_title="Object Detection App",
+    page_title="Object Detection",
     page_icon="🔍",
     layout="wide"
 )
 
-# Display environment info if needed for debugging
-if "debug" in st.experimental_get_query_params():
-    st.write("Python version:", sys.version)
-    st.write("Environment variables:", dict(os.environ))
-    st.write("Current working directory:", os.getcwd())
-    st.write("Directory contents:", os.listdir())
-
-# Create a sidebar
-st.sidebar.title("Object Detection App")
-st.sidebar.markdown("""
-This app uses Detectron2 to detect objects in images.
+# Display app header
+st.title("🔍 Object Detection with Detectron2")
+st.markdown("""
+Upload an image to detect objects using Facebook AI Research's Detectron2 model.
 """)
 
-# Display loading message
-with st.spinner("Loading dependencies..."):
-    try:
-        import torch
-        from detectron2.engine import DefaultPredictor
-        from detectron2.config import get_cfg
-        from detectron2 import model_zoo
-        from detectron2.utils.visualizer import Visualizer
-        from detectron2.data import MetadataCatalog
-        
-        st.sidebar.success("✅ Dependencies loaded successfully!")
-    except Exception as e:
-        st.error(f"Failed to load dependencies: {e}")
-        st.error(traceback.format_exc())
-        st.stop()
+# Setup sidebar
+with st.sidebar:
+    st.header("About")
+    st.markdown("This app uses Detectron2 to detect objects in images.")
+    
+    # Show environment info for debugging
+    if st.checkbox("Show Environment Info", False):
+        st.write("Python version:", sys.version)
+        st.write("Working directory:", os.getcwd())
+        st.write("Directory contents:", os.listdir())
+
+# Import Detectron2 with error handling
+try:
+    import torch
+    from detectron2 import model_zoo
+    from detectron2.engine import DefaultPredictor
+    from detectron2.config import get_cfg
+    from detectron2.utils.visualizer import Visualizer
+    from detectron2.data import MetadataCatalog
+except Exception as e:
+    st.error(f"Failed to import required libraries: {str(e)}")
+    st.error("Please check that all dependencies are correctly installed.")
+    st.stop()
 
 # Load the model
 @st.cache_resource
-def load_model():
+def load_detectron_model():
+    """Load the Detectron2 model with caching"""
     try:
-        # Configure the model
+        # Set up configuration
         cfg = get_cfg()
         cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
-        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Set threshold for object detection
         cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
+        cfg.MODEL.DEVICE = "cpu"  # Use CPU
         
-        # Use CPU for inference (more reliable in container environment)
-        cfg.MODEL.DEVICE = "cpu"
-        
-        # Initialize predictor
+        # Create predictor
         predictor = DefaultPredictor(cfg)
         return predictor, cfg
     except Exception as e:
-        st.error(f"Error loading model: {e}")
-        st.error(traceback.format_exc())
+        st.error(f"Error loading model: {str(e)}")
         return None, None
 
-# Main function
+# Process the image
+def process_image(image, predictor):
+    """Run object detection on the image"""
+    # Convert PIL Image to OpenCV format (RGB to BGR)
+    img = np.array(image.convert("RGB"))
+    
+    # Run inference
+    outputs = predictor(img)
+    
+    # Get the instances
+    instances = outputs["instances"].to("cpu")
+    
+    return img, instances
+
+# Visualize the results
+def visualize_results(img, instances, metadata):
+    """Create visualization of detection results"""
+    v = Visualizer(img, metadata=metadata, scale=1.2)
+    out = v.draw_instance_predictions(instances)
+    return out.get_image()
+
+# Main app logic
 def main():
-    st.title("🔍 Object Detection")
-    st.markdown("""
-    Upload an image to detect objects using Detectron2's Faster R-CNN model.
-    """)
+    # Load model with a spinner
+    with st.spinner("Loading model... (this may take a moment on first run)"):
+        predictor, cfg = load_detectron_model()
     
-    # Load model
-    with st.spinner("Loading model..."):
-        predictor, cfg = load_model()
-        
     if predictor is None:
-        st.error("Failed to load the model. Check the error messages.")
+        st.error("Failed to load the detection model.")
         return
     
-    # File uploader
-    uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
+    # Get metadata
+    metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "coco_2017_val")
     
+    # Create file uploader
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    
+    # Process the uploaded image
     if uploaded_file is not None:
         try:
-            # Read and display the image
+            # Display the original image
             image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-            
-            # Convert to numpy array
-            image_array = np.array(image.convert("RGB"))
+            st.image(image, caption="Uploaded Image", width=400)
             
-            # Perform inference
+            # Process the image
             with st.spinner("Detecting objects..."):
-                outputs = predictor(image_array)
-                
-                # Get instances
-                instances = outputs["instances"].to("cpu")
+                img, instances = process_image(image, predictor)
                 
-                # Create visualizer
-                v = Visualizer(image_array, 
-                               metadata=MetadataCatalog.get(cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "coco_2017_val"), 
-                               scale=1.2)
-                
-                # Draw predictions
-                result = v.draw_instance_predictions(instances)
-                result_image = result.get_image()
-                
-                # Display result
-                st.image(result_image, caption="Detection Result", use_column_width=True)
-                
-                # Show detection information
+                # Check if any objects were detected
                 if len(instances) > 0:
+                    # Visualize the results
+                    result_img = visualize_results(img, instances, metadata)
+                    
+                    # Display the result
+                    st.image(result_img, caption="Detection Result", width=800)
+                    
+                    # Show detection details
                     st.subheader(f"Detected {len(instances)} objects")
                     
-                    # Get class names
-                    metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "coco_2017_val")
-                    class_names = metadata.thing_classes
+                    class_names = metadata.get("thing_classes", None)
                     
-                    # Show detections
+                    # Display each detection
                     for i in range(len(instances)):
                         score = instances.scores[i].item()
                         class_id = instances.pred_classes[i].item()
-                        class_name = class_names[class_id]
-                        box = instances.pred_boxes[i].tensor.numpy()[0]
                         
-                        st.write(f"**{class_name}**: {score:.2f} confidence")
+                        if class_names:
+                            label = class_names[class_id]
+                        else:
+                            label = f"Class {class_id}"
+                            
+                        st.write(f"**{label}** (Confidence: {score:.2f})")
                 else:
                     st.info("No objects detected in this image.")
         
         except Exception as e:
-            st.error(f"Error processing image: {e}")
-            st.error(traceback.format_exc())
+            st.error(f"Error processing image: {str(e)}")
 
 if __name__ == "__main__":
     main()