Update plot_study() (ultralytics#2112)

Author: glenn-jocher

models/hub/yolov5l6.yaml

@@ -54,7 +54,7 @@ head:
 
     [ -1, 1, Conv, [ 768, 3, 2 ] ],
     [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
-    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P5/64-xlarge)
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
 
     [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
   ]

models/hub/yolov5m6.yaml

@@ -54,7 +54,7 @@ head:
 
     [ -1, 1, Conv, [ 768, 3, 2 ] ],
     [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
-    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P5/64-xlarge)
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
 
     [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
   ]

models/hub/yolov5s6.yaml

@@ -54,7 +54,7 @@ head:
 
     [ -1, 1, Conv, [ 768, 3, 2 ] ],
     [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
-    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P5/64-xlarge)
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
 
     [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
   ]

models/hub/yolov5x6.yaml

@@ -54,7 +54,7 @@ head:
 
     [ -1, 1, Conv, [ 768, 3, 2 ] ],
     [ [ -1, 12 ], 1, Concat, [ 1 ] ],  # cat head P6
-    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P5/64-xlarge)
+    [ -1, 3, C3, [ 1024, False ] ],  # 32 (P6/64-xlarge)
 
     [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ],  # Detect(P3, P4, P5, P6)
   ]

utils/plots.py

@@ -226,16 +226,17 @@ def plot_targets_txt():  # from utils.plots import *; plot_targets_txt()
 def plot_study_txt(path='', x=None):  # from utils.plots import *; plot_study_txt()
     # Plot study.txt generated by test.py
     fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
-    ax = ax.ravel()
+    # ax = ax.ravel()
 
     fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
-    for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]:
+    # for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]:
+    for f in sorted(Path(path).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', '[email protected]', '[email protected]:.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])
+        # 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
         ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8,
@@ -244,14 +245,14 @@ def plot_study_txt(path='', x=None):  # from utils.plots import *; plot_study_tx
     ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
              'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
 
-    ax2.grid()
-    ax2.set_yticks(np.arange(30, 60, 5))
+    ax2.grid(alpha=0.2)
+    ax2.set_yticks(np.arange(20, 60, 5))
     ax2.set_xlim(0, 30)
-    ax2.set_ylim(29, 51)
+    ax2.set_ylim(30, 55)
     ax2.set_xlabel('GPU Speed (ms/img)')
     ax2.set_ylabel('COCO AP val')
     ax2.legend(loc='lower right')
-    plt.savefig('test_study.png', dpi=300)
+    plt.savefig(str(Path(path).name) + '.png', dpi=300)
 
 
 def plot_labels(labels, save_dir=Path(''), loggers=None):