Plot image grid with labels.
(images,
batch_idx,
cls,
bboxes=np.zeros(0, dtype=np.float32),
masks=np.zeros(0, dtype=np.uint8),
kpts=np.zeros((0, 51), dtype=np.float32),
paths=None,
fname='images.jpg',
names=None,
on_plot=None)
| 381 | |
| 382 | @threaded |
| 383 | def plot_images(images, |
| 384 | batch_idx, |
| 385 | cls, |
| 386 | bboxes=np.zeros(0, dtype=np.float32), |
| 387 | masks=np.zeros(0, dtype=np.uint8), |
| 388 | kpts=np.zeros((0, 51), dtype=np.float32), |
| 389 | paths=None, |
| 390 | fname='images.jpg', |
| 391 | names=None, |
| 392 | on_plot=None): |
| 393 | """Plot image grid with labels.""" |
| 394 | if isinstance(images, torch.Tensor): |
| 395 | images = images.cpu().float().numpy() |
| 396 | if isinstance(cls, torch.Tensor): |
| 397 | cls = cls.cpu().numpy() |
| 398 | if isinstance(bboxes, torch.Tensor): |
| 399 | bboxes = bboxes.cpu().numpy() |
| 400 | if isinstance(masks, torch.Tensor): |
| 401 | masks = masks.cpu().numpy().astype(int) |
| 402 | if isinstance(kpts, torch.Tensor): |
| 403 | kpts = kpts.cpu().numpy() |
| 404 | if isinstance(batch_idx, torch.Tensor): |
| 405 | batch_idx = batch_idx.cpu().numpy() |
| 406 | |
| 407 | max_size = 1920 # max image size |
| 408 | max_subplots = 16 # max image subplots, i.e. 4x4 |
| 409 | bs, _, h, w = images.shape # batch size, _, height, width |
| 410 | bs = min(bs, max_subplots) # limit plot images |
| 411 | ns = np.ceil(bs ** 0.5) # number of subplots (square) |
| 412 | if np.max(images[0]) <= 1: |
| 413 | images *= 255 # de-normalise (optional) |
| 414 | |
| 415 | # Build Image |
| 416 | mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init |
| 417 | for i, im in enumerate(images): |
| 418 | if i == max_subplots: # if last batch has fewer images than we expect |
| 419 | break |
| 420 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin |
| 421 | im = im.transpose(1, 2, 0) |
| 422 | mosaic[y:y + h, x:x + w, :] = im |
| 423 | |
| 424 | # Resize (optional) |
| 425 | scale = max_size / ns / max(h, w) |
| 426 | if scale < 1: |
| 427 | h = math.ceil(scale * h) |
| 428 | w = math.ceil(scale * w) |
| 429 | mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h))) |
| 430 | |
| 431 | # Annotate |
| 432 | fs = int((h + w) * ns * 0.01) # font size |
| 433 | annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names) |
| 434 | for i in range(i + 1): |
| 435 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin |
| 436 | annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders |
| 437 | if paths: |
| 438 | annotator.text((x + 5, y + 5), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames |
| 439 | if len(cls) > 0: |
| 440 | idx = batch_idx == i |
no test coverage detected