Update app.py

app.py

@@ -0,0 +1,53 @@
+import os
+import torch
+from torchvision import transforms
+import torchvision.models.detection as detection
+import gradio as gr
+from PIL import Image
+import numpy as np
+import cv2
+
+# Load the trained model
+model = detection.fasterrcnn_resnet50_fpn(pretrained=False)
+num_classes = 91  # COCO has 80 classes + 1 background
+in_features = model.roi_heads.box_predictor.cls_score.in_features
+model.roi_heads.box_predictor = detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)
+
+# Load the model weights
+model.load_state_dict(torch.load('final_model.pth', weights_only=True))
+model.eval()
+
+# Define transformations
+transform = transforms.Compose([
+    transforms.Resize((600, 600)),
+    transforms.ToTensor(),
+])
+
+# Prediction function
+def predict(image):
+    image = transform(image).unsqueeze(0)  # Add batch dimension
+    with torch.no_grad():
+        predictions = model(image)
+    
+    # Process predictions
+    boxes = predictions[0]['boxes'].cpu().numpy()
+    scores = predictions[0]['scores'].cpu().numpy()
+    labels = predictions[0]['labels'].cpu().numpy()
+
+    # Filter out low-confidence predictions
+    threshold = 0.5
+    boxes = boxes[scores > threshold]
+    labels = labels[scores > threshold]
+
+    # Draw boxes on the image
+    image_np = np.array(image.squeeze().permute(1, 2, 0).cpu())
+    for box, label in zip(boxes, labels):
+        x1, y1, x2, y2 = box.astype(int)
+        image_np = cv2.rectangle(image_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
+        image_np = cv2.putText(image_np, str(label), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+
+    return Image.fromarray(image_np)
+
+# Gradio interface
+iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(type="pil"), outputs="image", title="Object Detection with Faster R-CNN")
+iface.launch()