changed the file app

app.py

@@ -1,71 +1,13 @@
 import streamlit as st
+from imageai.Detection import ObjectDetection
+import os
+import time
 from PIL import Image
-import torch
-from torchvision.models import detection
-from torchvision import transforms
-import numpy as np
-
-def load_model():
-    """Load the model directly from torchvision"""
-    with st.spinner('Loading model...'):
-        model = detection.retinanet_resnet50_fpn(pretrained=True)
-        model.eval()
-        return model
-
-def get_prediction(image, model, threshold=0.5):
-    """Get predictions for the image"""
-    # Transform the image
-    transform = transforms.Compose([
-        transforms.ToTensor()
-    ])
-    
-    img_tensor = transform(image)
-    
-    # Get prediction
-    with torch.no_grad():
-        prediction = model([img_tensor])
-    
-    # Get all the predicted class labels
-    pred_classes = [COCO_CLASSES[i] for i in prediction[0]['labels'].numpy()]
-    
-    # Get all the predicted bounding boxes
-    pred_boxes = prediction[0]['boxes'].numpy()
-    
-    # Get the predicted scores
-    pred_scores = prediction[0]['scores'].numpy()
-    
-    # Filter predictions based on threshold
-    mask = pred_scores >= threshold
-    boxes = pred_boxes[mask]
-    classes = np.array(pred_classes)[mask]
-    scores = pred_scores[mask]
-    
-    return boxes, classes, scores
-
-# COCO class labels
-COCO_CLASSES = [
-    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
-    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
-    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
-    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
-    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
-    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
-    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
-    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
-    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
-    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
-    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
-    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
-]
 
 def main():
     st.title("Object Detection App")
     st.write("Upload an image and get object detections!")
 
-    # Load model
-    model = load_model()
-    st.success("Model loaded successfully!")
-
     # File uploader
     uploaded_file = st.file_uploader("Choose an image...", type=['jpg', 'jpeg', 'png'])
 
@@ -73,36 +15,49 @@ def main():
         # Display the uploaded image
         image = Image.open(uploaded_file)
         st.image(image, caption='Uploaded Image', use_column_width=True)
+        
+        # Save the uploaded file temporarily
+        with open("temp_image.jpg", "wb") as f:
+            f.write(uploaded_file.getbuffer())
 
         # Button to start detection
         if st.button('Start Detection'):
-            with st.spinner('Performing detection...'):
+            # Spinner while model loads
+            with st.spinner('Loading model and performing detection...'):
                 try:
-                    # Get predictions
-                    boxes, classes, scores = get_prediction(image, model)
+                    execution_path = os.getcwd()
+                    detector = ObjectDetection()
+                    detector.setModelTypeAsRetinaNet()
+                    detector.setModelPath(os.path.join(execution_path, "retinanet_resnet50_fpn_coco-eeacb38b.pth"))
+                    detector.loadModel()
+
+                    # Perform detection
+                    detections = detector.detectObjectsFromImage(
+                        input_image="temp_image.jpg",
+                        output_image_path="output_image.jpg",
+                        minimum_percentage_probability=10
+                    )
 
                     # Display results
                     st.success('Detection Complete!')
                     
-                    # Create a nice looking table for detections
-                    st.write("### Detected Objects:")
+                    # Display detected image
+                    detected_image = Image.open("output_image.jpg")
+                    st.image(detected_image, caption='Detected Objects', use_column_width=True)
                     
-                    # Create columns for better visualization
-                    col1, col2 = st.columns(2)
-                    
-                    with col1:
-                        st.write("Object")
-                    with col2:
-                        st.write("Confidence")
-                    
-                    for cls, score in zip(classes, scores):
-                        with col1:
-                            st.write(f"{cls}")
-                        with col2:
-                            st.write(f"{score*100:.2f}%")
+                    # Display detections with probabilities
+                    st.write("### Detected Objects:")
+                    for obj in detections:
+                        st.write(f"- {obj['name']}: {obj['percentage_probability']:.2f}%")
 
                 except Exception as e:
                     st.error(f"An error occurred: {str(e)}")
 
+            # Clean up temporary files
+            if os.path.exists("temp_image.jpg"):
+                os.remove("temp_image.jpg")
+            if os.path.exists("output_image.jpg"):
+                os.remove("output_image.jpg")
+
 if __name__ == "__main__":
     main()