Update app.py

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
 import torch
-import torchvision.transforms.functional as F
-from torchvision.models.detection import fasterrcnn_resnet50_fpn 
 import numpy as np
+from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
 from PIL import Image, ImageDraw
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
-from torchvision import datasets, models
 
 # Initialize Faster R-CNN model
-model = models.detection.fasterrcnn_resnet50_fpn_v2(pretrained=True)
+model = fasterrcnn_resnet50_fpn_v2(pretrained=True)
 
 # Load the saved state_dict
 state_dict = torch.load("faster-rcnn.pth", map_location=torch.device('cpu'))
@@ -36,11 +34,15 @@ def draw_boxes_with_labels(image, boxes, labels):
 
     for box, label in zip(boxes, labels):
         x1, y1, x2, y2 = box
+
+        # Correct the box coordinates if they are not integers
+        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+        
         color = COLOR_MAP[label]
         draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
 
         text_position = (x1, y1 - 10)
-        
+
         # Draw label background
         text_bbox = draw.textbbox(text_position, label)
         draw.rectangle(text_bbox, fill=color)
@@ -92,9 +94,6 @@ def predict_image(image):
     # Convert the original image back to PIL format
     pil_image = Image.fromarray(image)
 
-    # Scale image_tensor back to uint8 format
-    image_uint8 = (image_tensor * 255).byte().cpu()
-
     # Draw the boxes and labels on the image
     drawn_image = draw_boxes_with_labels(pil_image, boxes.cpu(), label_names)