Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,195 +1,136 @@
 import streamlit as st
-import torch
 import numpy as np
 import cv2
 import os
-import tempfile
+import sys
 from PIL import Image
 import traceback
-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
 
-# Setup page config
+# Configure the app
 st.set_page_config(
-    page_title="Object Detection with Detectron2",
+    page_title="Object Detection App",
     page_icon="🔍",
     layout="wide"
 )
 
-# Print environment info for debugging
+# 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())
-    st.write("Temp directory:", tempfile.gettempdir())
-    st.write("PyTorch CUDA available:", torch.cuda.is_available())
 
-# Load the Detectron2 model
-@st.cache_resource
-def load_model():
-    try:
-        # Set the configuration for the 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  # Set threshold for detection confidence
-        cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
-        
-        # Use GPU if available, otherwise use CPU
-        if torch.cuda.is_available():
-            st.sidebar.success("GPU is available! Using CUDA.")
-            cfg.MODEL.DEVICE = "cuda"
-        else:
-            st.sidebar.info("GPU not available. Using CPU.")
-            cfg.MODEL.DEVICE = "cpu"
-            
-        # Initialize the predictor
-        predictor = DefaultPredictor(cfg)
-        return predictor
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        st.error(traceback.format_exc())
-        return None
+# Create a sidebar
+st.sidebar.title("Object Detection App")
+st.sidebar.markdown("""
+This app uses Detectron2 to detect objects in images.
+""")
 
-# Function for image prediction
-def predict_fn(image, predictor):
+# Display loading message
+with st.spinner("Loading dependencies..."):
     try:
-        # Convert the PIL image to a format the model can use
-        image_array = np.array(image.convert("RGB"))
-        
-        # Make predictions
-        outputs = predictor(image_array)
-        
-        # Get the predicted classes and bounding boxes
-        instances = outputs["instances"].to("cpu")
-        pred_classes = instances.pred_classes.numpy()
-        pred_boxes = instances.pred_boxes.tensor.numpy()
-        scores = instances.scores.numpy()
-        
-        # Get class names from metadata
-        metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "coco_2017_val")
-        class_names = metadata.thing_classes
+        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
         
-        return pred_classes, pred_boxes, scores, image_array, class_names
+        st.sidebar.success("✅ Dependencies loaded successfully!")
     except Exception as e:
-        st.error(f"Error during prediction: {e}")
+        st.error(f"Failed to load dependencies: {e}")
         st.error(traceback.format_exc())
-        return None, None, None, np.array(image), None
+        st.stop()
 
-# Function to visualize predictions
-def visualize_predictions(image_array, instances, class_names):
+# Load the model
+@st.cache_resource
+def load_model():
     try:
-        # Create a visualizer object
-        metadata = MetadataCatalog.get("coco_2017_val")
-        if class_names is not None:
-            metadata.thing_classes = class_names
-            
-        v = Visualizer(image_array, metadata=metadata, scale=1.2)
+        # Configure the 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")
         
-        # Draw the predictions on the image
-        vis_output = v.draw_instance_predictions(instances.to("cpu"))
+        # Use CPU for inference (more reliable in container environment)
+        cfg.MODEL.DEVICE = "cpu"
         
-        # Return the visualized image
-        return vis_output.get_image()
+        # Initialize predictor
+        predictor = DefaultPredictor(cfg)
+        return predictor, cfg
     except Exception as e:
-        st.error(f"Error during visualization: {e}")
+        st.error(f"Error loading model: {e}")
         st.error(traceback.format_exc())
-        return image_array
+        return None, None
 
-# Main application
+# Main function
 def main():
-    st.title("🔍 Object Detection with Detectron2")
+    st.title("🔍 Object Detection")
     st.markdown("""
-    Upload an image to detect objects using Facebook AI Research's Detectron2 model.
-    This demo uses the Faster R-CNN model with ResNet-50-FPN backbone.
+    Upload an image to detect objects using Detectron2's Faster R-CNN model.
     """)
     
-    # Create a sidebar
-    st.sidebar.title("Settings")
-    confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.05)
-    
-    # Load the model
-    with st.spinner("Loading model... This might take a minute on first run."):
-        predictor = load_model()
-    
+    # Load model
+    with st.spinner("Loading model..."):
+        predictor, cfg = load_model()
+        
     if predictor is None:
-        st.error("Failed to load model. Please check the error messages above.")
+        st.error("Failed to load the model. Check the error messages.")
         return
     
-    # Update model confidence threshold
-    predictor.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
-    
-    # Upload image
-    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-    
-    col1, col2 = st.columns(2)
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
     
     if uploaded_file is not None:
-        # Display original image
-        image = Image.open(uploaded_file)
-        with col1:
-            st.subheader("Original Image")
-            st.image(image, use_column_width=True)
-        
-        # Process the image
-        with st.spinner("Detecting objects..."):
-            try:
-                # Make predictions
-                pred_classes, pred_boxes, scores, image_array, class_names = predict_fn(image, predictor)
+        try:
+            # Read and display the 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"))
+            
+            # Perform inference
+            with st.spinner("Detecting objects..."):
+                outputs = predictor(image_array)
                 
-                if pred_classes is not None and len(pred_classes) > 0:
-                    # Reconstruct instances for visualization
-                    from detectron2.structures import Instances
-                    h, w = image_array.shape[:2]
-                    instances = Instances((h, w))
-                    instances.pred_boxes = torch.tensor(pred_boxes)
-                    instances.scores = torch.tensor(scores)
-                    instances.pred_classes = torch.tensor(pred_classes)
-                    
-                    # Visualize predictions
-                    result_image = visualize_predictions(image_array, instances, class_names)
-                    
-                    # Display results
-                    with col2:
-                        st.subheader("Detected Objects")
-                        st.image(result_image, use_column_width=True)
-                    
-                    # Display detection details
-                    st.subheader("Detection Results")
+                # Get instances
+                instances = outputs["instances"].to("cpu")
+                
+                # 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
+                if len(instances) > 0:
+                    st.subheader(f"Detected {len(instances)} objects")
                     
-                    # Create a table for results
-                    results_data = []
-                    for i, (cls, box, score) in enumerate(zip(pred_classes, pred_boxes, scores)):
-                        class_name = class_names[cls] if class_names is not None else f"Class {cls}"
-                        results_data.append({
-                            "Object": class_name,
-                            "Confidence": f"{score:.2f}",
-                            "Bounding Box": f"[{int(box[0])}, {int(box[1])}, {int(box[2])}, {int(box[3])}]"
-                        })
+                    # Get class names
+                    metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "coco_2017_val")
+                    class_names = metadata.thing_classes
                     
-                    # Display results as table
-                    if results_data:
-                        st.table(results_data)
-                    else:
-                        st.info("No objects detected.")
+                    # Show detections
+                    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")
                 else:
-                    with col2:
-                        st.info("No objects detected in the image.")
-            except Exception as e:
-                st.error(f"Error processing image: {e}")
-                st.error(traceback.format_exc())
-    else:
-        # Show a sample image option
-        if st.button("Use a sample image"):
-            # You could include a sample image in your repository
-            # and load it here to demonstrate the functionality
-            st.info("Sample image option selected - this would load a demo image if implemented")
-    
-    # Footer
-    st.markdown("---")
-    st.markdown("Built with Streamlit and Facebook AI Research's Detectron2")
+                    st.info("No objects detected in this image.")
+        
+        except Exception as e:
+            st.error(f"Error processing image: {e}")
+            st.error(traceback.format_exc())
 
 if __name__ == "__main__":
     main()