From 4f44aaf26bb10134699cc65ed5dd2ae6792e8e12 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 30 May 2020 15:18:59 -0700
Subject: [PATCH] updates

---
 utils/utils.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/utils/utils.py b/utils/utils.py
index c65d762..c85a9a7 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -471,6 +471,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T
     # Settings
     merge = True  # merge for best mAP
     min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
+    max_det = 300  # maximum number of detections per image
     time_limit = 10.0  # seconds to quit after
 
     t = time.time()
@@ -520,6 +521,8 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T
         c = x[:, 5] * 0 if agnostic else x[:, 5]  # classes
         boxes, scores = x[:, :4].clone() + c.view(-1, 1) * max_wh, x[:, 4]  # boxes (offset by class), scores
         i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
+        if i.shape[0] > max_det:  # limit detections
+            i = i[:max_det]
         if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
             try:  # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
                 iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
@@ -975,15 +978,27 @@ def plot_targets_txt():  # from utils.utils import *; plot_targets_txt()
 
 def plot_study_txt(f='study.txt', x=None):  # from utils.utils import *; plot_study_txt()
     # Plot study.txt generated by test.py
-    y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
-    x = np.arange(y.shape[1]) if x is None else np.array(x)
-    s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
     fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
     ax = ax.ravel()
-    for i in range(7):
-        ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
-        ax[i].set_title(s[i])
-    plt.savefig(f.replace('.txt','.png'), dpi=200)
+
+    for f in glob.glob('study*.txt'):
+        y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
+        x = np.arange(y.shape[1]) if x is None else np.array(x)
+        s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
+        for i in range(7):
+            ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
+            ax[i].set_title(s[i])
+
+        j = y[3].argmax() + 1
+        ax[7].plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8, label=Path(f).stem)
+        ax[7].plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [33.5, 39.1, 42.5, 45.9, 49., 50.5],
+                   '.-', linewidth=2, markersize=8, label='EfficientDet')
+        ax[7].set_xlabel('Latency (ms)')
+        ax[7].set_ylabel('COCO AP val')
+
+    ax[7].legend()
+    ax[7].set_xlim(0)
+    plt.savefig(f.replace('.txt', '.png'), dpi=200)
 
 
 def plot_labels(labels):