Spaces:
Runtime error
Runtime error
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() | |