This commit is contained in:
Glenn Jocher 2019-07-09 18:16:15 +02:00
parent 9c227dd2b4
commit 5b0ba6d7b2
1 changed files with 2 additions and 1 deletions

View File

@ -616,6 +616,7 @@ def plot_images(imgs, targets, paths=None, fname='images.jpg'):
# Plots training images overlaid with targets # Plots training images overlaid with targets
imgs = imgs.cpu().numpy() imgs = imgs.cpu().numpy()
targets = targets.cpu().numpy() targets = targets.cpu().numpy()
# targets = targets[targets[:, 1] == 21] # plot only one class
fig = plt.figure(figsize=(10, 10)) fig = plt.figure(figsize=(10, 10))
bs, _, h, w = imgs.shape # batch size, _, height, width bs, _, h, w = imgs.shape # batch size, _, height, width
@ -626,7 +627,7 @@ def plot_images(imgs, targets, paths=None, fname='images.jpg'):
boxes[[0, 2]] *= w boxes[[0, 2]] *= w
boxes[[1, 3]] *= h boxes[[1, 3]] *= h
plt.subplot(ns, ns, i + 1).imshow(imgs[i].transpose(1, 2, 0)) plt.subplot(ns, ns, i + 1).imshow(imgs[i].transpose(1, 2, 0))
plt.plot(boxes[[0, 2, 2, 0, 0]], boxes[[1, 1, 3, 3, 1]], '.-') plt.plot(boxes[[0, 2, 2, 0, 0]], boxes[[1, 1, 3, 3, 1]], 'w.-')
plt.axis('off') plt.axis('off')
if paths is not None: if paths is not None:
s = Path(paths[i]).name s = Path(paths[i]).name