Spaces:
Sleeping
Sleeping
| import contextlib | |
| import math | |
| from pathlib import Path | |
| import cv2 | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from .. import threaded | |
| from ..general import xywh2xyxy | |
| from ..plots import Annotator, colors | |
| def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg', names=None): | |
| # Plot image grid with labels | |
| if isinstance(images, torch.Tensor): | |
| images = images.cpu().float().numpy() | |
| if isinstance(targets, torch.Tensor): | |
| targets = targets.cpu().numpy() | |
| if isinstance(masks, torch.Tensor): | |
| masks = masks.cpu().numpy().astype(int) | |
| max_size = 1920 # max image size | |
| max_subplots = 16 # max image subplots, i.e. 4x4 | |
| bs, _, h, w = images.shape # batch size, _, height, width | |
| bs = min(bs, max_subplots) # limit plot images | |
| ns = np.ceil(bs ** 0.5) # number of subplots (square) | |
| if np.max(images[0]) <= 1: | |
| images *= 255 # de-normalise (optional) | |
| # Build Image | |
| mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init | |
| for i, im in enumerate(images): | |
| if i == max_subplots: # if last batch has fewer images than we expect | |
| break | |
| x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin | |
| im = im.transpose(1, 2, 0) | |
| mosaic[y:y + h, x:x + w, :] = im | |
| # Resize (optional) | |
| scale = max_size / ns / max(h, w) | |
| if scale < 1: | |
| h = math.ceil(scale * h) | |
| w = math.ceil(scale * w) | |
| mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h))) | |
| # Annotate | |
| fs = int((h + w) * ns * 0.01) # font size | |
| annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names) | |
| for i in range(i + 1): | |
| x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin | |
| annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders | |
| if paths: | |
| annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames | |
| if len(targets) > 0: | |
| idx = targets[:, 0] == i | |
| ti = targets[idx] # image targets | |
| boxes = xywh2xyxy(ti[:, 2:6]).T | |
| classes = ti[:, 1].astype('int') | |
| labels = ti.shape[1] == 6 # labels if no conf column | |
| conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred) | |
| if boxes.shape[1]: | |
| if boxes.max() <= 1.01: # if normalized with tolerance 0.01 | |
| boxes[[0, 2]] *= w # scale to pixels | |
| boxes[[1, 3]] *= h | |
| elif scale < 1: # absolute coords need scale if image scales | |
| boxes *= scale | |
| boxes[[0, 2]] += x | |
| boxes[[1, 3]] += y | |
| for j, box in enumerate(boxes.T.tolist()): | |
| cls = classes[j] | |
| color = colors(cls) | |
| cls = names[cls] if names else cls | |
| if labels or conf[j] > 0.25: # 0.25 conf thresh | |
| label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}' | |
| annotator.box_label(box, label, color=color) | |
| # Plot masks | |
| if len(masks): | |
| if masks.max() > 1.0: # mean that masks are overlap | |
| image_masks = masks[[i]] # (1, 640, 640) | |
| nl = len(ti) | |
| index = np.arange(nl).reshape(nl, 1, 1) + 1 | |
| image_masks = np.repeat(image_masks, nl, axis=0) | |
| image_masks = np.where(image_masks == index, 1.0, 0.0) | |
| else: | |
| image_masks = masks[idx] | |
| im = np.asarray(annotator.im).copy() | |
| for j, box in enumerate(boxes.T.tolist()): | |
| if labels or conf[j] > 0.25: # 0.25 conf thresh | |
| color = colors(classes[j]) | |
| mh, mw = image_masks[j].shape | |
| if mh != h or mw != w: | |
| mask = image_masks[j].astype(np.uint8) | |
| mask = cv2.resize(mask, (w, h)) | |
| mask = mask.astype(bool) | |
| else: | |
| mask = image_masks[j].astype(bool) | |
| with contextlib.suppress(Exception): | |
| im[y:y + h, x:x + w, :][mask] = im[y:y + h, x:x + w, :][mask] * 0.4 + np.array(color) * 0.6 | |
| annotator.fromarray(im) | |
| annotator.im.save(fname) # save | |
| def plot_results_with_masks(file='path/to/results.csv', dir='', best=True): | |
| # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv') | |
| save_dir = Path(file).parent if file else Path(dir) | |
| fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True) | |
| ax = ax.ravel() | |
| files = list(save_dir.glob('results*.csv')) | |
| assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.' | |
| for f in files: | |
| try: | |
| data = pd.read_csv(f) | |
| index = np.argmax(0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] + | |
| 0.1 * data.values[:, 11]) | |
| s = [x.strip() for x in data.columns] | |
| x = data.values[:, 0] | |
| for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]): | |
| y = data.values[:, j] | |
| # y[y == 0] = np.nan # don't show zero values | |
| ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=2) | |
| if best: | |
| # best | |
| ax[i].scatter(index, y[index], color='r', label=f'best:{index}', marker='*', linewidth=3) | |
| ax[i].set_title(s[j] + f'\n{round(y[index], 5)}') | |
| else: | |
| # last | |
| ax[i].scatter(x[-1], y[-1], color='r', label='last', marker='*', linewidth=3) | |
| ax[i].set_title(s[j] + f'\n{round(y[-1], 5)}') | |
| # if j in [8, 9, 10]: # share train and val loss y axes | |
| # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) | |
| except Exception as e: | |
| print(f'Warning: Plotting error for {f}: {e}') | |
| ax[1].legend() | |
| fig.savefig(save_dir / 'results.png', dpi=200) | |
| plt.close() | |