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 + h), | |
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() | |