Update plots.py

Error running python detect.py --visualize ultralytics#5503
ys31jp · Nov 5, 2021 · 56535ba · 56535ba
1 parent 99a45ba
commit 56535ba
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/utils/plots.py b/utils/plots.py
@@ -443,5 +443,7 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
                 ax[i].axis('off')
 
             print(f'Saving {save_dir / f}... ({n}/{channels})')
-            plt.savefig(save_dir / f, dpi=300, bbox_inches='tight')
+            if not os.path.exists(save_dir):
+                os.makedirs(save_dir)
+            plt.savefig('{}/{}.png'.format(save_dir, f), dpi=300, bbox_inches='tight')
             plt.close()