mirror of https://github.com/WongKinYiu/yolov7.git
183 lines
7.7 KiB
Python
183 lines
7.7 KiB
Python
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_masks(img, masks, colors, alpha=0.5):
|
|
"""
|
|
Args:
|
|
img (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
|
|
masks (tensor): predicted masks on cuda, shape: [n, h, w]
|
|
colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
|
|
Return:
|
|
ndarray: img after draw masks, shape: [h, w, 3]
|
|
|
|
transform colors and send img_gpu to cpu for the most time.
|
|
"""
|
|
img_gpu = img.clone()
|
|
num_masks = len(masks)
|
|
if num_masks == 0:
|
|
return img.permute(1, 2, 0).contiguous().cpu().numpy() * 255
|
|
|
|
# [n, 1, 1, 3]
|
|
# faster this way to transform colors
|
|
colors = torch.tensor(colors, device=img.device).float() / 255.0
|
|
colors = colors[:, None, None, :]
|
|
# [n, h, w, 1]
|
|
masks = masks[:, :, :, None]
|
|
masks_color = masks.repeat(1, 1, 1, 3) * colors * alpha
|
|
inv_alph_masks = masks * (-alpha) + 1
|
|
masks_color_summand = masks_color[0]
|
|
if num_masks > 1:
|
|
inv_alph_cumul = inv_alph_masks[:(num_masks - 1)].cumprod(dim=0)
|
|
masks_color_cumul = masks_color[1:] * inv_alph_cumul
|
|
masks_color_summand += masks_color_cumul.sum(dim=0)
|
|
|
|
# print(inv_alph_masks.prod(dim=0).shape) # [h, w, 1]
|
|
img_gpu = img_gpu.flip(dims=[0]) # filp channel for opencv
|
|
img_gpu = img_gpu.permute(1, 2, 0).contiguous()
|
|
# [h, w, 3]
|
|
img_gpu = img_gpu * inv_alph_masks.prod(dim=0) + masks_color_summand
|
|
return (img_gpu * 255).byte().cpu().numpy()
|
|
|
|
|
|
@threaded
|
|
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(np.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()
|