Plot training labels including class histograms and box statistics.
(boxes, cls, names=(), save_dir=Path(''), on_plot=None)
| 278 | @TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395 |
| 279 | @plt_settings() |
| 280 | def plot_labels(boxes, cls, names=(), save_dir=Path(''), on_plot=None): |
| 281 | """Plot training labels including class histograms and box statistics.""" |
| 282 | import pandas as pd |
| 283 | import seaborn as sn |
| 284 | |
| 285 | # Filter matplotlib>=3.7.2 warning and Seaborn use_inf and is_categorical FutureWarnings |
| 286 | warnings.filterwarnings('ignore', category=UserWarning, message='The figure layout has changed to tight') |
| 287 | warnings.filterwarnings('ignore', category=FutureWarning) |
| 288 | |
| 289 | # Plot dataset labels |
| 290 | LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ") |
| 291 | nc = int(cls.max() + 1) # number of classes |
| 292 | boxes = boxes[:1000000] # limit to 1M boxes |
| 293 | x = pd.DataFrame(boxes, columns=['x', 'y', 'width', 'height']) |
| 294 | |
| 295 | # Seaborn correlogram |
| 296 | sn.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9)) |
| 297 | plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200) |
| 298 | plt.close() |
| 299 | |
| 300 | # Matplotlib labels |
| 301 | ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel() |
| 302 | y = ax[0].hist(cls, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8) |
| 303 | for i in range(nc): |
| 304 | y[2].patches[i].set_color([x / 255 for x in colors(i)]) |
| 305 | ax[0].set_ylabel('instances') |
| 306 | if 0 < len(names) < 30: |
| 307 | ax[0].set_xticks(range(len(names))) |
| 308 | ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10) |
| 309 | else: |
| 310 | ax[0].set_xlabel('classes') |
| 311 | sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9) |
| 312 | sn.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9) |
| 313 | |
| 314 | # Rectangles |
| 315 | boxes[:, 0:2] = 0.5 # center |
| 316 | boxes = ops.xywh2xyxy(boxes) * 1000 |
| 317 | img = Image.fromarray(np.ones((1000, 1000, 3), dtype=np.uint8) * 255) |
| 318 | for cls, box in zip(cls[:500], boxes[:500]): |
| 319 | ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot |
| 320 | ax[1].imshow(img) |
| 321 | ax[1].axis('off') |
| 322 | |
| 323 | for a in [0, 1, 2, 3]: |
| 324 | for s in ['top', 'right', 'left', 'bottom']: |
| 325 | ax[a].spines[s].set_visible(False) |
| 326 | |
| 327 | fname = save_dir / 'labels.jpg' |
| 328 | plt.savefig(fname, dpi=200) |
| 329 | plt.close() |
| 330 | if on_plot: |
| 331 | on_plot(fname) |
| 332 | |
| 333 | |
| 334 | def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True): |