YOLOv5s输出合并特征图并保持大小一致

在utils/plots.py中添加以下代码（只有这一段与参考链接中不同，其他一样），剩下的步骤可参考【yolov5】改进系列——特征图可视化

import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import torch

def feature_visualization_all(x, module_type, stage, save_dir=Path('runs/detect/exp')):
    """
    x:              Features to be visualized
    module_type:    Module type
    stage:          Module stage within model
    save_dir:       Directory to save results
    """
    if 'Detect' not in module_type:
        batch, channels, height, width = x.shape  # batch, channels, height, width
        if height > 1 and width > 1:
            f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png"  # filename
            
            # Sum along the channel dimension and transpose to get proper orientation
            #img = x[0].cpu().sum(0).detach().numpy()
            img = x[0].cpu().transpose(0, 1).sum(1).detach().numpy()

            # Set a fixed figure size to ensure consistent output size
            plt.figure(figsize=(6, 6))
            plt.imshow(img)
            plt.axis('off')  # Turn off axes for clean visualization

            # Save the image with specified dpi and no extra margins
            plt.savefig(f, bbox_inches='tight')
            plt.close()

            LOGGER.info(f'Saving {f}...'