Modified app

app.py

@@ -1,8 +1,57 @@
+from typing import List
+import cv2
+import torch
 import numpy as np
 import gradio as gr
-from detection import detect_objects
-from config import PASCAL_CLASSES
+import config as modelConfig
 
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+from yolov3 import YOLOv3
+import utils
+from utils import cells_to_bboxes, non_max_suppression, draw_bounding_boxes, YoloGradCAM
+
+
+model = YOLOv3(num_classes=len(modelConfig.PASCAL_CLASSES))
+optimizer = torch.optim.Adam(model.parameters(), lr=0.00072/100, weight_decay=1e-4)
+
+utils.load_checkpoint("checkpoint.pth.tar", model, optimizer, 0.00072/100)
+
+scaled_anchors = (
+    torch.tensor(modelConfig.ANCHORS)
+    * torch.tensor(modelConfig.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
+).to(modelConfig.DEVICE)
+
+yolo_grad_cam = YoloGradCAM(model=model, target_layers=[model.layers[-2]], use_cuda=False)
+
+@torch.inference_mode()
+def detect_objects(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.4, enable_grad_cam: bool = False, transparency: float = 0.5) -> List[np.ndarray]:
+    transformed_image = modelConfig.transforms(image=image)["image"].unsqueeze(0)
+    transformed_image = transformed_image.cuda()
+    output = model(transformed_image)
+
+    bboxes = [[] for _ in range(1)]
+    for i in range(3):
+        batch_size, A, S, _, _ = output[i].shape
+        anchor = scaled_anchors[i]
+        boxes_scale_i = cells_to_bboxes(
+            output[i], anchor, S=S, is_preds=True
+        )
+        for idx, (box) in enumerate(boxes_scale_i):
+            bboxes[idx] += box
+
+    nms_boxes = non_max_suppression(
+        bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
+    )
+    plot_img_with_bboxes = draw_bounding_boxes(image.copy(), nms_boxes, class_labels=modelConfig.PASCAL_CLASSES)
+    if not enable_grad_cam:
+        return [plot_img_with_bboxes]
+
+    grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :]
+    img = cv2.resize(image, (416, 416))
+    img = np.float32(img) / 255
+    grad_cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True, image_weight=transparency)
+    return [plot_img_with_bboxes, grad_cam_image]
 
 def inference(
     image: np.ndarray,
@@ -10,18 +59,22 @@ def inference(
     enable_grad_cam: str,
     transparency: float,
 ):
-    infer_output = detect_objects(image, iou_thresh, thresh, enable_grad_cam, transparency)
-    return infer_output
-
-
-title = "YoloV3 for Pascal VOC Dataset"
-description = f"Pytorch Implementation of YoloV3 model trained on Pascal VOC dataset with GradCAM \n Classes in pascol voc are: {', '.join(PASCAL_CLASSES)}"
-example_images = [
-    ["images/001114.jpg", 0.7, 0.5, True, 0.6],
-    ["images/001133.jpg", 0.6, 0.5, True, 0.6],
-    ["images/001142.jpg", 0.65, 0.45, True, 0.6],
-    ["images/001147.jpg", 0.6, 0.5, True, 0.6],
-    ["images/001155.jpg", 0.7, 0.7, True, 0.6],
+    results = detect_objects(image, iou_thresh, thresh, enable_grad_cam, transparency)
+    return results
+
+
+title = "Object detection application using YoloV3 Model"
+description = f"Object detection application using pre-trained YoloV3 model for Pascal VOC dataset. This app has GradCAM option also. \n Classes in Pascal voc dataset are : {', '.join(modelConfig.PASCAL_CLASSES)}"
+examples = [
+    ["images_2/000811.jpg", 0.6, 0.6, True, 0.6],
+    ["images_2/000830.jpg", 0.5, 0.5, True, 0.6],
+    ["images_2/000842.jpg", 0.6, 0.6, True, 0.6],    
+    ["images_2/001114.jpg", 0.4, 0.5, True, 0.6],
+    ["images_2/001133.jpg", 0.7, 0.7, True, 0.6],
+    ["images_2/001155.jpg", 0.7, 0.69, True, 0.6],
+    ["images_2/000008.jpg", 0.66, 0.69, True, 0.6],
+    ["images_2/000031.jpg", 0.6, 0.6, True, 0.6],
+    ["images_2/000175.jpg", 0.6, 0.6, True, 0.6],
 ]
 
 demo = gr.Interface(
@@ -38,7 +91,7 @@ demo = gr.Interface(
     ],
     title=title,
     description=description,
-    examples=example_images,
+    examples=examples,
 )
 
 demo.launch(debug=True)