Skip to content

Commit

Permalink
ConfusionMatrix normalize=True fix (ultralytics#3587)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Jun 11, 2021
1 parent ec2da4a commit e8c5237
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions utils/metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,9 +161,8 @@ def matrix(self):
def plot(self, normalize=True, save_dir='', names=()):
try:
import seaborn as sn

if normalize:
array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6) # normalize columns

array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-6) if normalize else 1) # normalize columns
array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)

fig = plt.figure(figsize=(12, 9), tight_layout=True)
Expand All @@ -178,7 +177,7 @@ def plot(self, normalize=True, save_dir='', names=()):
fig.axes[0].set_ylabel('Predicted')
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
except Exception as e:
pass
print(f'WARNING: ConfusionMatrix plot failure: {e}')

def print(self):
for i in range(self.nc + 1):
Expand Down

0 comments on commit e8c5237

Please sign in to comment.