Update app.py

app.py

@@ -1,19 +1,19 @@
 import gradio as gr
-import cv2
 import torch
-from torchvision.models.detection import fasterrcnn_resnet50_fpn
-import numpy as np
+from torchvision.utils import draw_bounding_boxes
+import matplotlib.pyplot as plt
+from PIL import Image, ImageDraw
 import torchvision.transforms.functional as F
+import cv2
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
-import matplotlib.pyplot as plt
-
-# Define class names and color mapping
-CLASS_NAMES = ['bone-fracture', 'angle', 'fracture', 'line', 'messed_up_angle']
 
-def map_label_to_class(label_id):
-    return CLASS_NAMES[label_id - 1]
+# Load your trained model
+model = torch.load("faster_rcnn.pth", map_location=torch.device('cpu'))
+model.eval()
 
+# Define class names and color map
+CLASS_NAMES = ['bone-fracture', 'angle', 'fracture', 'line', 'messed_up_angle']
 COLOR_MAP = {
     'bone-fracture': (255, 0, 0),
     'angle': (0, 255, 0),
@@ -22,96 +22,75 @@ COLOR_MAP = {
     'messed_up_angle': (255, 165, 0)
 }
 
-def draw_boxes_with_labels(image, boxes, labels):
-    # Convert the image from tensor to numpy
-    image_np = image.permute(1, 2, 0).numpy() * 255
-    image_np = image_np.astype(np.uint8)
+def map_label_to_class(label_id):
+    return CLASS_NAMES[label_id - 1]
 
-    # Plot the image
-    plt.figure(figsize=(10, 10))
-    plt.imshow(image_np)
+def draw_boxes_with_labels(image, boxes, labels):
+    pil_image = F.to_pil_image(image)
+    draw = ImageDraw.Draw(pil_image)
 
-    # Draw bounding boxes and labels
     for box, label in zip(boxes, labels):
-        x1, y1, x2, y2 = box.int().tolist()
+        x1, y1, x2, y2 = box
         color = COLOR_MAP[label]
+        draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
+        text_position = (x1, y1 - 10)
+        text_bbox = draw.textbbox(text_position, label)
+        draw.rectangle(text_bbox, fill=color)
+        draw.text(text_position, label, fill="white")
 
-        # Draw rectangle
-        plt.gca().add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor=color, facecolor='none'))
-
-        # Draw label
-        plt.text(x1, y1, label, color='white', fontsize=12, bbox=dict(facecolor=color, alpha=0.5))
-
-    plt.axis('off')
-    plt.close()  # Prevents the figure from showing in an interactive mode
-    return image_np  # Return the image as a numpy array
+    return F.to_tensor(pil_image)
 
-def get_transforms(train):
+def get_transforms(train=False):
     if train:
-        return A.Compose([
-            A.HorizontalFlip(p=0.5),
-            A.RandomBrightnessContrast(p=0.2),
-            A.Resize(height=800, width=800),
-            ToTensorV2(),
-        ])
+        transform = A.Compose([
+            A.Resize(800, 800), 
+            A.HorizontalFlip(p=0.3),
+            A.VerticalFlip(p=0.3),
+            A.RandomBrightnessContrast(p=0.1),
+            A.ColorJitter(p=0.1),
+            ToTensorV2()
+        ], bbox_params=A.BboxParams(format='coco'))
     else:
-        return A.Compose([
-            A.Resize(height=800, width=800),
-            ToTensorV2(),
-        ])
-
-def predict_image(image, model, device, transform, threshold=0.1):
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        transform = A.Compose([
+            A.Resize(800, 800),
+            ToTensorV2()
+        ], bbox_params=A.BboxParams(format='coco'))
+    return transform
+
+def predict_image(image):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Preprocess the input image
+    image = np.array(image)
+    transform = get_transforms(train=False)
     transformed = transform(image=image)
-    image_tensor = transformed['image'].to(device)
+    image_tensor = transformed['image'].unsqueeze(0).to(device)
     image_tensor = image_tensor.float() / 255.0
 
-    model.eval()
+    # Predict using the model
     with torch.no_grad():
-        prediction = model([image_tensor])
-        pred = prediction[0]
-
-    boxes = pred['boxes'][pred['scores'] > threshold]
-    labels = pred['labels'][pred['scores'] > threshold]
-
-    # Convert the image to uint8 format for drawing
-    image_uint8 = (image_tensor * 255).byte().cpu()
-
-    # Map labels to class names
+        prediction = model(image_tensor)
+    
+    # Extract predictions (boxes, labels, scores)
+    pred = prediction[0]
+    boxes = pred['boxes'][pred['scores'] > 0.5]
+    labels = pred['labels'][pred['scores'] > 0.5]
+    scores = pred['scores'][pred['scores'] > 0.5]
+    
+    # Convert labels to class names
     label_names = [map_label_to_class(label.item()) for label in labels]
 
-    # Draw bounding boxes with labels
+    # Convert image to uint8 for visualization
+    image_uint8 = (image_tensor.squeeze(0) * 255).byte().cpu()
+    
+    # Draw the boxes and labels on the image
     drawn_image = draw_boxes_with_labels(image_uint8, boxes.cpu(), label_names)
-    return drawn_image  # Return the modified image
-# Tải mô hình
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = fasterrcnn_resnet50_fpn(pretrained=False)
-
-pretrained_dict = torch.load('faster-rcnn.pth', map_location=torch.device('cpu'))
-model_dict = model.state_dict()
-# Lọc các khóa không cần thiết
-pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
-# Ghi đè các mục trong state_dict hiện có
-model_dict.update(pretrained_dict)
-# Tải state_dict mới
-model.load_state_dict(model_dict)
-model.to(device).eval()
+    
+    return F.to_pil_image(drawn_image)
 
-# Define Gradio interface function
-def process_image(input_image):
-    transform = get_transforms(train=False)
-    output_image = predict_image(input_image, model, device, transform)
+# Define the Gradio interface
+def predict(img):
+    output_image = predict_image(img)
     return output_image
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="numpy"),
-    title="Object Detection with Faster R-CNN",
-    description="Upload an image to detect objects."
-)
-
-# Run the application
-if __name__ == "__main__":
-    iface.launch(share=True)  # To create a public link
+gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Image()).launch()