Create app.py

app.py

@@ -0,0 +1,120 @@
+import streamlit as st
+import torch
+import torchvision.transforms as T
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+# Load a pre-trained Faster R-CNN model
+@st.cache_resource
+def load_model():
+    model = fasterrcnn_resnet50_fpn(pretrained=True)
+    model.eval()
+    return model
+
+model = load_model()
+
+# Replace this in your code
+OFFICE_INSTANCE_CATEGORY_NAMES = [
+    '__background__',  # Always include the background class
+    'person',
+    'chair',
+    'couch',
+    'potted plant',
+    'dining table',
+    'tv',
+    'laptop',
+    'mouse',
+    'remote',
+    'keyboard',
+    'cell phone',
+    'book',
+    'clock',
+    'vase',
+    'scissors',
+    'teddy bear',
+    'hair drier',
+    'toothbrush',
+    'bottle',
+    'cup',
+    'fork',
+    'knife',
+    'spoon',
+    'bowl',
+    'backpack',
+    'handbag',
+    'tie',
+    'suitcase',
+    'frisbee',
+    'sports ball',
+    'surfboard',
+    'tennis racket',
+    'wine glass',
+    'sandwich',
+    'orange',
+    'pizza',
+    'donut',
+    'cake',
+    'microwave',
+    'oven',
+    'toaster',
+    'sink',
+    'refrigerator',
+    'apple',
+    'clock',
+    'book',
+    'vase',
+    'bench',
+    'umbrella'
+]
+
+# Define a transformation to convert the image to tensor
+transform = T.Compose([
+    T.ToTensor(),
+])
+
+# Function to draw bounding boxes on the image
+def draw_boxes(image, predictions, threshold=0.5):
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    
+    for idx, box in enumerate(predictions[0]['boxes']):
+        score = predictions[0]['scores'][idx].item()
+        label = COCO_INSTANCE_CATEGORY_NAMES[predictions[0]['labels'][idx]]
+        
+        if score >= threshold:
+            box = box.tolist()
+            draw.rectangle(box, outline="red", width=3)
+            text = f"{label}: {int(score * 100)}%"
+            text_size = font.getsize(text)
+            draw.rectangle([box[0], box[1] - text_size[1], box[0] + text_size[0], box[1]], fill="red")
+            draw.text((box[0], box[1] - text_size[1]), text, fill="white", font=font)
+
+# Streamlit interface
+def main():
+    st.title("Object Detection with Faster R-CNN")
+    st.write("Upload an image and detect objects using Faster R-CNN.")
+
+    # File uploader
+    uploaded_image = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
+
+    if uploaded_image is not None:
+        image = Image.open(uploaded_image).convert("RGB")
+        st.image(image, caption="Uploaded Image", use_container_width=True)
+
+        if st.button("Analyze"):
+            with st.spinner("Analyzing..."):
+                image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+
+                # Perform object detection
+                with torch.no_grad():
+                    predictions = model(image_tensor)
+
+                # Draw bounding boxes
+                image_with_boxes = image.copy()
+                draw_boxes(image_with_boxes, predictions)
+
+                st.image(image_with_boxes, caption="Detected Objects", use_container_width=True)
+
+if __name__ == "__main__":
+    main()