Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from model import Lightning_YOLO
 import config
-from utils import non_max_suppression, cells_to_bboxes, draw_bounding_boxes
+from utils import non_max_suppression, cells_to_bboxes, draw_bounding_boxes, get_annotations
 import torch
 
 scaled_anchors = config.scaled_anchors
@@ -11,9 +11,11 @@ model.load_state_dict(torch.load("yolov3.pth", map_location="cpu"), strict=False
 model.eval()
 
 def inference(image, threst = 0.5, iou_tresh = 0.5):
+    image_copy = image.copy()
     transformed_image = config.transforms(image=image)["image"].unsqueeze(0)
     output = model(transformed_image)
     bboxes = [[] for _ in range(1)]
+    nms_boxes_output = []
     for i in range(3):
       batch_size, A, S, _, _ = output[i].shape
       anchor = scaled_anchors[i]
@@ -23,12 +25,22 @@ def inference(image, threst = 0.5, iou_tresh = 0.5):
       for idx, (box) in enumerate(boxes_scale_i):
           bboxes[idx] += box
 
-    nms_boxes = non_max_suppression(
-      bboxes[0], iou_threshold=iou_tresh, threshold=threst, box_format="midpoint",
-    )
-    plot_img = draw_bounding_boxes(image.copy(), nms_boxes, class_labels=config.PASCAL_CLASSES)
+    
+    # nms_boxes = non_max_suppression(
+    #   bboxes[0], iou_threshold=iou_tresh, threshold=threst, box_format="midpoint",
+    # )
+    for i in range(image.shape[0]):
+        
+        nms_boxes = non_max_suppression(
+          bboxes[i], iou_threshold=iou_tresh, threshold=threst, box_format="midpoint",
+        )
+        nms_boxes_output.append(nms_boxes)
+    
+    annotations = get_annotations(nms_boxes_output,config.IMAGE_SIZE,config.IMAGE_SIZE)
+    # plot_img = draw_bounding_boxes(image.copy(), nms_boxes, class_labels=config.PASCAL_CLASSES)
 
-    return plot_img
+    # return plot_img
+    return [image_copy, annotations]
 
 def visualize_gradcam(image, target_layer=-5, show_cam=True, transparency=0.5):
   # if show_cam:
@@ -67,7 +79,10 @@ window2 = gr.Interface(
         gr.Slider(0, 1, value=0.5, step=0.1, label="Transparency", info="Set Transparency of GRAD-CAMs"), 
     ],
     outputs=[
-        gr.Image(label="Grad-CAM Visualization"),
+        # gr.Image(label="Grad-CAM Visualization"),
+        gr.AnnotatedImage(label='BBox Prediction',
+                        height=config.IMAGE_SIZE,
+                        width=config.IMAGE_SIZE)
     ],
     # examples=ex2,
 )