Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,65 +6,79 @@ from PIL import Image
 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
 
-# Function to load the model
-def load_model(model_path):
-    # Initialize the configuration
+# Load the Detectron2 model
+@st.cache_resource
+def load_model():
+    # Set the configuration for the model (COCO pre-trained model for detection)
     cfg = get_cfg()
     cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
-    cfg.MODEL.WEIGHTS = model_path  # Path to your trained model
-    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Threshold for prediction
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Set the threshold for prediction
+    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  # Model weights from the model zoo
+    cfg.MODEL.DEVICE = "cpu"  # Change to "cuda" if using GPU
+
+    # Initialize the predictor
     predictor = DefaultPredictor(cfg)
     return predictor
 
-# Function to make predictions
-def predict_fn(image, model_path):
-    # Convert image to numpy array
-    image = np.array(image)
-    
-    # Load model
-    predictor = load_model(model_path)
-    
+# Load the model
+predictor = load_model()
+
+# Function for image prediction
+def predict_fn(image):
+    # Convert the PIL image to a format the model can use
+    # Convert to RGB (Detectron2 uses BGR format internally)
+    image = np.array(image.convert("RGB"))
+
     # Make predictions
     outputs = predictor(image)
+
+    # 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()
+
+    return pred_classes, pred_boxes, image
+
+# Function to display image with bounding boxes
+def visualize_predictions(image, pred_classes, pred_boxes):
+    # Create a visualizer object
+    v = Visualizer(image[:, :, ::-1], MetadataCatalog.get("coco_2017_val"), scale=1.2)
+    v = v.draw_instance_predictions(pred_classes)
     
-    # Extract predicted bounding boxes and class labels
-    instances = outputs["instances"]
-    pred_classes = instances.pred_classes.to("cpu").numpy()
-    pred_boxes = instances.pred_boxes.tensor.to("cpu").numpy()
-    
-    return pred_classes, pred_boxes
+    # Draw the bounding boxes on the image
+    for box in pred_boxes:
+        start_point = tuple(map(int, box[:2]))
+        end_point = tuple(map(int, box[2:]))
+        color = (0, 255, 0)  # Green color
+        thickness = 2
+        image = cv2.rectangle(image, start_point, end_point, color, thickness)
+
+    return image
 
-# Streamlit UI elements
-st.title("Wildlife Detection with Your Model")
+# Streamlit UI
+st.title("Object Detection with Detectron2")
 
-# Upload an image for testing
-uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+# Upload image
+uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 
 if uploaded_image is not None:
-    # Open the uploaded image using PIL
+    # Open the uploaded image
     image = Image.open(uploaded_image)
     
-    # Display the uploaded image
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+    # Make predictions
+    pred_classes, pred_boxes, image_array = predict_fn(image)
     
-    # Set model path
-    model_path = "sandbox338/wild-life-model"  # Your model path on Hugging Face
+    # Visualize predictions on the image
+    image_with_boxes = visualize_predictions(image_array, pred_classes, pred_boxes)
     
-    # Make predictions
-    pred_classes, pred_boxes = predict_fn(image, model_path)
+    # Convert the image back to RGB format for display in Streamlit
+    image_with_boxes = Image.fromarray(image_with_boxes[:, :, ::-1])  # Convert BGR to RGB
     
-    # Display predictions
-    st.write(f"Predicted classes: {pred_classes}")
-    st.write(f"Predicted bounding boxes: {pred_boxes}")
+    # Display the processed image with bounding boxes
+    st.image(image_with_boxes, caption="Processed Image", use_column_width=True)
 
-    # Optionally, visualize predictions on the image
-    img_array = np.array(image)
-    for i, box in enumerate(pred_boxes):
-        start_point = (int(box[0]), int(box[1]))
-        end_point = (int(box[2]), int(box[3]))
-        img_array = cv2.rectangle(img_array, start_point, end_point, (0, 255, 0), 2)
-    
-    st.image(img_array, caption="Predicted Image with Bounding Boxes", use_column_width=True)
-else:
-    st.write("Please upload an image to detect wildlife!")
+    # Display the classes detected
+    st.write("Predicted Classes:", pred_classes)