Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,143 +1,53 @@
 import streamlit as st
-import numpy as np
 from PIL import Image
-import cv2
-import os
-import sys
-
-# Set page configuration
-st.set_page_config(
-    page_title="Object Detection",
-    page_icon="🔍",
-    layout="wide"
-)
-
-# Display app header
-st.title("🔍 Object Detection with Detectron2")
-st.markdown("""
-Upload an image to detect objects using Facebook AI Research's Detectron2 model.
-""")
-
-# 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
+import numpy as np
+import torch
+import asyncio
+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
+
+# Fix event loop issue
 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_detectron_model():
-    """Load the Detectron2 model with caching"""
-    try:
-        # 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  # 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
-        
-        # Create predictor
-        predictor = DefaultPredictor(cfg)
-        return predictor, cfg
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None, None
+    asyncio.get_running_loop()
+except RuntimeError:
+    asyncio.set_event_loop(asyncio.new_event_loop())
 
-# 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
+# Title and uploader
+st.title("Detectron2 Object Detection")
+st.write("Upload an image to perform object detection")
 
-# 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()
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-# Main app logic
-def main():
-    # Load model with a spinner
-    with st.spinner("Loading model... (this may take a moment on first run)"):
-        predictor, cfg = load_detectron_model()
-    
-    if predictor is None:
-        st.error("Failed to load the detection model.")
-        return
-    
-    # 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:
-            # Display the original image
-            image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image", width=400)
-            
-            # Process the image
-            with st.spinner("Detecting objects..."):
-                img, instances = process_image(image, predictor)
-                
-                # 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")
-                    
-                    class_names = metadata.get("thing_classes", None)
-                    
-                    # Display each detection
-                    for i in range(len(instances)):
-                        score = instances.scores[i].item()
-                        class_id = instances.pred_classes[i].item()
-                        
-                        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: {str(e)}")
-
-if __name__ == "__main__":
-    main()
+@st.cache_resource
+def load_model():
+    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.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
+    predictor = DefaultPredictor(cfg)
+    return predictor
+
+def predict_fn(predictor, image):
+    image_array = np.array(image)[:, :, :3]
+    outputs = predictor(image_array)
+    return outputs["instances"], image_array
+
+def visualize_predictions(image, instances):
+    v = Visualizer(image[:, :, ::-1], MetadataCatalog.get("coco_2017_val"), scale=1.2)
+    v = v.draw_instance_predictions(instances)
+    result = v.get_image()
+    return result[:, :, ::-1]
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+    st.write("Processing...")
+
+    predictor = load_model()
+    instances, image_array = predict_fn(predictor, image)
+    result_image = visualize_predictions(image_array, instances)
+
+    st.image(result_image, caption="Detected Objects", use_column_width=True)