This commit is contained in:
Glenn Jocher 2019-04-09 13:21:39 +02:00
parent 3e85a4191a
commit d8cbf9b7a7
1 changed files with 5 additions and 2 deletions

View File

@ -489,14 +489,17 @@ def plot_wh_methods(): # from utils.utils import *; plot_wh_methods()
def plot_images(imgs, targets, fname='images.jpg'):
# Plots training images overlaid with targets
imgs = imgs.cpu().numpy()
targets = targets.cpu().numpy()
fig = plt.figure(figsize=(10, 10))
img_size = imgs.shape[3]
bs = imgs.shape[0] # batch size
sp = np.ceil(bs ** 0.5) # subplots
for i in range(bs):
boxes = xywh2xyxy(targets[targets[:, 0] == i, 2:6]).numpy().T * img_size
plt.subplot(sp, sp, i + 1).imshow(imgs[i].cpu().numpy().transpose(1, 2, 0))
boxes = xywh2xyxy(targets[targets[:, 0] == i, 2:6]).T * img_size
plt.subplot(sp, sp, i + 1).imshow(imgs[i].transpose(1, 2, 0))
plt.plot(boxes[[0, 2, 2, 0, 0]], boxes[[1, 1, 3, 3, 1]], '.-')
plt.axis('off')
fig.tight_layout()