extra

app.py

@@ -16,9 +16,9 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 import numpy as np
 from albumentations.pytorch import ToTensorV2
 
-from pytorchyolo.utils.utils import load_classes, rescale_boxes, non_max_suppression, print_environment_info
+from pytorchyolo.utils.utils import load_classes, rescale_boxes, print_environment_info
 
-from utils_for_app import cells_to_bboxes,plot_image,YoloCAM
+from utils_for_app import cells_to_bboxes,plot_image,YoloCAM,non_max_suppression
 from yolov3 import YOLOv3
 from loss import YoloLoss
 from utils import LearningRateFinder
@@ -81,7 +81,7 @@ def process_image_and_plot(image,iou_threshold=0.5, threshold=0.4):
     with torch.no_grad(): 
         detections = model(transformed_image)
         detect = non_max_suppression(detections, iou_threshold, threshold)
-        BBOX = rescale_boxes(detect[0], img_size, image.shape[:2])
+        # BBOX = rescale_boxes(detect[0], img_size, image.shape[:2])
     
 
     fig = plot_image(transformed_image[0].permute(1, 2, 0), detect[0])

pytorchyolo/utils/utils.py

@@ -310,6 +310,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
     """
 
     nc = prediction.shape[2] - 5  # number of classes
+    
+    print("Class Count", nc)
 
     # Settings
     # (pixels) minimum and maximum box width and height

utils_for_app.py

@@ -100,7 +100,7 @@ def intersection_over_union(boxes_preds, boxes_labels, box_format="midpoint"):
 
     return intersection / (box1_area + box2_area - intersection + 1e-6)
 
-def non_max_suppressionv2(bboxes, iou_threshold, threshold, box_format="corners"):
+def non_max_suppression(bboxes, iou_threshold, threshold, box_format="corners"):
     """
     Video explanation of this function:
     https://youtu.be/YDkjWEN8jNA
@@ -118,8 +118,8 @@ def non_max_suppressionv2(bboxes, iou_threshold, threshold, box_format="corners"
     
     assert type(bboxes) == list
 
-    bboxes = [box for box in bboxes if box[1] > threshold]
-    bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
+    bboxes = [box for box in bboxes if box[4] > threshold]
+    bboxes = sorted(bboxes, key=lambda x: x[4], reverse=True)
     bboxes_after_nms = []
 
     while bboxes:
@@ -130,8 +130,8 @@ def non_max_suppressionv2(bboxes, iou_threshold, threshold, box_format="corners"
             for box in bboxes
             if box[0] != chosen_box[0]
             or intersection_over_union(
-                torch.tensor(chosen_box[2:]),
-                torch.tensor(box[2:]),
+                torch.tensor(chosen_box[0:4]),
+                torch.tensor(box[0:4]),
                 box_format=box_format,
             )
             < iou_threshold
@@ -144,6 +144,8 @@ def non_max_suppressionv2(bboxes, iou_threshold, threshold, box_format="corners"
 
     return bboxes_after_nms
 
+
+
 def plot_image(image, boxes):
     """Plots predicted bounding boxes on the image"""
     cmap = plt.get_cmap("tab20b")
@@ -164,7 +166,7 @@ def plot_image(image, boxes):
     # for x1, y1, x2, y2, conf, cls_pred in detections:
     for box in boxes:
         assert len(box) == 6, "box should contain class pred, confidence, x, y, width, height"
-        class_pred = box[-1]
+        class_pred = box[5]
         # box = box[2:]
         upper_left_x = box[0] - box[2] / 2
         upper_left_y = box[1] - box[3] / 2