diff --git a/utils/plots.py b/utils/plots.py index 7af3dc5..a9b5870 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -1,13 +1,13 @@ # Plotting utils import glob -import math import os import random from copy import copy from pathlib import Path import cv2 +import math import matplotlib import matplotlib.pyplot as plt import numpy as np @@ -218,13 +218,13 @@ def plot_targets_txt(): # from utils.plots import *; plot_targets_txt() plt.savefig('targets.jpg', dpi=200) -def plot_study_txt(f='study.txt', x=None): # from utils.plots import *; plot_study_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() fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True) - for f in ['study/study_coco_%s.txt' % x for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]: + for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]: 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)'] @@ -234,7 +234,7 @@ def plot_study_txt(f='study.txt', x=None): # from utils.plots import *; plot_st j = y[3].argmax() + 1 ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8, - label=Path(f).stem.replace('study_coco_', '').replace('yolo', 'YOLO')) + label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO')) 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') @@ -246,8 +246,7 @@ def plot_study_txt(f='study.txt', x=None): # from utils.plots import *; plot_st ax2.set_xlabel('GPU Speed (ms/img)') ax2.set_ylabel('COCO AP val') ax2.legend(loc='lower right') - plt.savefig('study_mAP_latency.png', dpi=300) - plt.savefig(f.replace('.txt', '.png'), dpi=300) + plt.savefig('test_study.png', dpi=300) def plot_labels(labels, save_dir=''):