asdTools-ReID热力图可视化

文章首发见博客：https://mwhls.top/4869.html。
无图/格式错误/后续更新请见首发页。
更多更新请到mwhls.top查看
欢迎留言提问或批评建议，私信不回。

Github - 开源代码及Readme
Blog - 工具介绍

摘要：基于TorchCam实现ReID的热力图可视化的工具介绍。

场景

• 想为科研添砖加瓦，因此需要一个好看的可视化。
• 或者只是想了解一下模型的效果。

使用示例

• 开源地址：asdTools/asdTools/Tools/Image/VisualizeHeatmapOfReID.py at main
• 处理结果示例：asdTools/Sample/VisualizeHeatmapOfReID at main
• 实例化一个VisualizeHeatmapOfReID对象，将 模型、图片文件夹、测试时transform 传入即可。
  • 使用示例见if __name__ == "__main__":
• 安装 torch-cam: pip install torchcam or conda install -c frgfm torchcam, GitHub: https://github.com/frgfm/torch-cam

    from torchvision.transforms.functional import to_pil_imagefrom asdTools.Classes.Image.ImageBase import ImageBasefrom torchcam.utils import overlay_maskfrom torchcam.methods import GradCAMclass VisualizeHeatmapOfReID(ImageBase):""" Sample: Sample/VisualizeHeatmapOfReID使用torch-cam可视化热力图，仅测试于ReID模型。Visualize heatmap by torch-cam, test only on ReID.torch-cam: `pip install torchcam` or `conda install -c frgfm torchcam`, GitHub: https://github.com/frgfm/torch-cam"""def __init__(self, **kwargs) -> None:super().__init__(multipleFiles=True, **kwargs)def __call__(self, model, imgs_dir:str, transform,img_ext:list=["png", "jpg", "jpeg"],torchCamMethod=GradCAM,device="cuda") -> str:self.run(model, imgs_dir, transform, img_ext, torchCamMethod, device)def run(self, model, imgs_dir:str, transform, img_ext:list, torchCamMethod, device) -> str:self.begining()# get paths of imgs from img_dirimgs_path = self.get_paths_from_dir(imgs_dir)self.log(f"{len(imgs_path)} files found in {imgs_dir}")imgs_path = self.filter_ext(imgs_path, img_ext)self.log(f"{len(imgs_path)} images found after filter extension by {img_ext}")# init torchCammodel.to(device).eval()cam_extractor = torchCamMethod(model)for i, img_path in enumerate(imgs_path):# model(x)img = self.read_img(img_path)x = transform(img).unsqueeze(0).to(device)out = model(x)# visualize featureactivation_map = cam_extractor(class_idx=0, scores=out.unsqueeze(0))[0]result = overlay_mask(img, to_pil_image(activation_map, mode='F'), alpha=0.5)# save imgsave_path = self.remove_root_of_path(path=img_path, root=imgs_dir)save_middle_dir = self.get_dir_of_file(save_path)save_name = self.get_name_of_file(save_path, True)save_path = self.save_image(result, output_middle_dir=save_middle_dir, output_file=save_name)self.log(f"{i+1}/{len(imgs_path)}: the heatmap of {img_path} has been saved to {save_path}.")self.done()if __name__ == "__main__":""" ---------- usage example ---------- -- you can simply insert VisualizeHeatmapOfReID() in the train.py or main.py-- and then set value of imgs_dirfrom asdTools.Tools.Image.VisualizeHeatmapOfReID import VisualizeHeatmapOfReIDfrom models.img_resnet import ResNet50import data.img_transforms as Timgs_dir = "../../../../Datasets/PRCC"weight = "logs/0/baseline.pth.tar"checkpoint = torch.load(weight)model = ResNet50(config)model.load_state_dict(checkpoint['model_state_dict'])model.cuda().eval()transform_test = T.Compose([T.Resize((256, 128)),T.ToTensor(),T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])color_heatmap = VisualizeHeatmapOfReID()color_heatmap(model, imgs_dir, transform_test)"""