Skip to content

Commit 248504c

Browse files
authored
Feature visualization improvements 32 (#3947)
1 parent dabad57 commit 248504c

2 files changed

Lines changed: 7 additions & 6 deletions

File tree

detect.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ def run(weights='yolov5s.pt', # model.pt path(s)
103103
t1 = time_synchronized()
104104
pred = model(img,
105105
augment=augment,
106-
visualize=increment_path(save_dir / 'features', mkdir=True) if visualize else False)[0]
106+
visualize=increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False)[0]
107107

108108
# Apply NMS
109109
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

utils/plots.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import yaml
1717
from PIL import Image, ImageDraw, ImageFont
1818

19-
from utils.general import increment_path, xywh2xyxy, xyxy2xywh
19+
from utils.general import xywh2xyxy, xyxy2xywh
2020
from utils.metrics import fitness
2121

2222
# Settings
@@ -447,7 +447,7 @@ def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''):
447447
fig.savefig(Path(save_dir) / 'results.png', dpi=200)
448448

449449

450-
def feature_visualization(x, module_type, stage, n=64, save_dir=Path('runs/detect/exp')):
450+
def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')):
451451
"""
452452
x: Features to be visualized
453453
module_type: Module type
@@ -460,13 +460,14 @@ def feature_visualization(x, module_type, stage, n=64, save_dir=Path('runs/detec
460460
if height > 1 and width > 1:
461461
f = f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
462462

463-
plt.figure(tight_layout=True)
464463
blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels
465464
n = min(n, channels) # number of plots
466-
ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True)[1].ravel() # 8 rows x n/8 cols
465+
fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols
466+
ax = ax.ravel()
467+
plt.subplots_adjust(wspace=0.05, hspace=0.05)
467468
for i in range(n):
468469
ax[i].imshow(blocks[i].squeeze()) # cmap='gray'
469470
ax[i].axis('off')
470471

471472
print(f'Saving {save_dir / f}... ({n}/{channels})')
472-
plt.savefig(save_dir / f, dpi=300)
473+
plt.savefig(save_dir / f, dpi=300, bbox_inches='tight')

0 commit comments

Comments
 (0)