Upload utils.py

utils.py

@@ -275,6 +275,7 @@ def plot_image(image, boxes):
         )
 
     plt.show()
+    return plt
 
 
 def get_evaluation_bboxes(
@@ -521,6 +522,34 @@ def plot_couple_examples(model, loader, thresh, iou_thresh, anchors):
             bboxes[i], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
         )
         plot_image(x[i].permute(1,2,0).detach().cpu(), nms_boxes)
+        
+def plot_single_image(model, image, thresh, iou_thresh, anchors):
+    print('calling to plot single image')
+    model.eval()
+    with torch.no_grad():
+        x = image
+        print('Train begin')
+        out = model(x)
+        print('Train end')
+        bboxes = [[] for _ in range(x.shape[0])]
+        for i in range(3):
+            batch_size, A, S, _, _ = out[i].shape
+            anchor = anchors[i]
+            boxes_scale_i = cells_to_bboxes(
+                out[i], anchor, S=S, is_preds=True
+            )
+            for idx, (box) in enumerate(boxes_scale_i):
+                bboxes[idx] += box
+
+        model.train()
+    print('batch size = ',batch_size)
+    for i in range(batch_size):
+        nms_boxes = non_max_suppression(
+            bboxes[i], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
+        )
+        print('calling to plot bounding boxes image')
+        fig = plot_image(x[i].permute(1,2,0).detach().cpu(), nms_boxes) 
+        return fig