src/streamlit_app.py

import streamlit as st
import torch
import numpy as np
import cv2
import os
import tempfile
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
st.set_page_config(
    page_title="Object Detection with Detectron2",
    page_icon="🔍",
    layout="wide"
)

# Print environment info for debugging
if "debug" in st.experimental_get_query_params():
    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

# Function for image prediction
def predict_fn(image, predictor):
    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
        
        return pred_classes, pred_boxes, scores, image_array, class_names
    except Exception as e:
        st.error(f"Error during prediction: {e}")
        st.error(traceback.format_exc())
        return None, None, None, np.array(image), None

# Function to visualize predictions
def visualize_predictions(image_array, instances, class_names):
    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)
        
        # Draw the predictions on the image
        vis_output = v.draw_instance_predictions(instances.to("cpu"))
        
        # Return the visualized image
        return vis_output.get_image()
    except Exception as e:
        st.error(f"Error during visualization: {e}")
        st.error(traceback.format_exc())
        return image_array

# Main application
def main():
    st.title("🔍 Object Detection with Detectron2")
    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.
    """)
    
    # 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()
    
    if predictor is None:
        st.error("Failed to load model. Please check the error messages above.")
        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)
    
    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)
                
                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")
                    
                    # 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])}]"
                        })
                    
                    # Display results as table
                    if results_data:
                        st.table(results_data)
                    else:
                        st.info("No objects detected.")
                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")

if __name__ == "__main__":
    main()