def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')): """ x: Features to be visualized module_type: Module type stage: Module stage within model n: Maximum number of feature maps to plot 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 blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels n = min(n, channels) # number of plots fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols ax = ax.ravel() plt.subplots_adjust(wspace=0.05, hspace=0.05) for i in range(n): ax[i].imshow(blocks[i].squeeze()) # cmap='gray' ax[i].axis('off') LOGGER.info(f'Saving {f}... ({n}/{channels})') plt.savefig(f, dpi=300, bbox_inches='tight') plt.close() np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy()) # npy save

这段代码是用于可视化特征图的函数。它接受参数x（要可视化的特征）、module_type（模块类型）、stage（模型中的模块阶段）、n（要绘制的最大特征图数量）和save_dir（保存结果的目录）。

函数首先检查module_type中是否包含"Detect"，如果不包含，则假设x的形状为(batch, channels, height, width)。然后根据height和width的值确定特征图是否是二维的。

接下来，函数将特征图拆分为多个块，每个块对应一个通道。然后，它创建一个子图网格，将每个特征图块绘制在其中，并保存结果。

最后，函数使用numpy保存特征图的数组，并返回保存的文件路径。

请注意，该代码依赖于torch、matplotlib和numpy等库。你需要确保这些库已经安装在你的环境中才能运行该代码。