使用Grad-CAM方法建立YOLOv5中实现三种不同的热力图可视化普通热力图、仅标签框内的热力图、去除红色以外区域的热力图文章目录使用Grad-CAM方法建立YOLOv5中实现三种不同的热力图可视化普通热力图、仅标签框内的热力图、去除红色以外区域的热力图步骤 1: 导入必要的库步骤 2: 加载YOLOv5模型步骤 3: 定义Grad-CAM类步骤 4: 处理输入图像步骤 5: 获取热力图特定区域热力图仅标签框内的热力图去除红色以外区域的热力图完整代码YOLOv5热力图grad-cam建立目标3种热力图可视化:普通的仅标签框内的去除红色以外区域的用于YOLOv5添加注意力机制、改进模型等YOLOv5中实现三种不同的热力图可视化普通热力图、仅标签框内的热力图、去除红色以外区域的热力图我们可以使用Grad-CAM方法。代码示例仅供参考。如何加载模型、处理图像、生成热力图以及展示结果。步骤 1: 导入必要的库importcv2importnumpyasnpimporttorchfromtorchvision.transformsimportfunctionalasFfrompytorch_grad_camimportGradCAM,ScoreCAM,GradCAMPlusPlus,AblationCAM,XGradCAM,EigenCAM,EigenGradCAM,LayerCAMfrompytorch_grad_cam.utils.imageimportshow_cam_on_imagefrommodels.experimentalimportattempt_load# 假设这是加载YOLOv5模型的方式步骤 2: 加载YOLOv5模型modelattempt_load(yolov5_weights.pt,map_locationcpu)# 或 cuda 如果有GPU的话model.eval()步骤 3: 定义Grad-CAM类target_layermodel.model[-1]# 这里假设目标层是最后一个层camGradCAM(modelmodel,target_layertarget_layer,use_cudaTrueiftorch.cuda.is_available()elseFalse)步骤 4: 处理输入图像defpreprocess_image(img_path):imgcv2.imread(img_path)imgcv2.cvtColor(img,cv2.COLOR_BGR2RGB)img_tensorF.to_tensor(img).unsqueeze(0)# 添加batch维度returnimg,img_tensor步骤 5: 获取热力图defget_gradcam_visualization(model,img_tensor,cam_method,target_categoryNone):grayscale_camcam_method(input_tensorimg_tensor,target_categorytarget_category)grayscale_camgrayscale_cam[0,:]visualizationshow_cam_on_image(img/255.0,grayscale_cam,use_rgbTrue)returnvisualization# 普通热力图img,img_tensorpreprocess_image(path_to_your_image.jpg)visualizationget_gradcam_visualization(model,img_tensor,cam)# 显示结果importmatplotlib.pyplotasplt plt.imshow(visualization)plt.show()特定区域热力图仅标签框内的热力图defget_bounding_boxes(model,img_tensor):withtorch.no_grad():outputmodel(img_tensor)boxesoutput.xyxy[0].tolist()returnboxesdefvisualize_bounding_box_heatmap(img,boxes,visualization):forboxinboxes:x1,y1,x2,y2int(box[0]),int(box[1]),int(box[2]),int(box[3])cropped_imgimg[y1:y2,x1:x2]cropped_tensorF.to_tensor(cropped_img).unsqueeze(0)cropped_visualizationget_gradcam_visualization(model,cropped_tensor,cam)img[y1:y2,x1:x2]cropped_visualizationreturnimg boxesget_bounding_boxes(model,img_tensor)heatmap_with_boxesvisualize_bounding_box_heatmap(img,boxes,visualization)plt.imshow(heatmap_with_boxes)plt.show()去除红色以外区域的热力图defremove_non_red_regions(visualization):lower_rednp.array([200,0,0],dtypeuint8)upper_rednp.array([255,100,100],dtypeuint8)maskcv2.inRange(visualization,lower_red,upper_red)visualization[mask0][255,255,255]# 将非红色区域变白或其他颜色returnvisualization heatmap_without_non_redremove_non_red_regions(visualization)plt.imshow(heatmap_without_non_red)plt.show()完整代码importcv2importnumpyasnpimporttorchfromtorchvision.transformsimportfunctionalasFfrompytorch_grad_camimportGradCAM,ScoreCAM,GradCAMPlusPlus,AblationCAM,XGradCAM,EigenCAM,EigenGradCAM,LayerCAMfrompytorch_grad_cam.utils.imageimportshow_cam_on_imagefrommodels.experimentalimportattempt_load# 加载YOLOv5模型modelattempt_load(yolov5_weights.pt,map_locationcpu)model.eval()# 定义Grad-CAM类target_layermodel.model[-1]camGradCAM(modelmodel,target_layertarget_layer,use_cudaTrueiftorch.cuda.is_available()elseFalse)# 处理输入图像defpreprocess_image(img_path):imgcv2.imread(img_path)imgcv2.cvtColor(img,cv2.COLOR_BGR2RGB)img_tensorF.to_tensor(img).unsqueeze(0)returnimg,img_tensor# 获取热力图defget_gradcam_visualization(model,img_tensor,cam_method,target_categoryNone):grayscale_camcam_method(input_tensorimg_tensor,target_categorytarget_category)grayscale_camgrayscale_cam[0,:]visualizationshow_cam_on_image(img/255.0,grayscale_cam,use_rgbTrue)returnvisualization# 获取边界框defget_bounding_boxes(model,img_tensor):withtorch.no_grad():outputmodel(img_tensor)boxesoutput.xyxy[0].tolist()returnboxes# 可视化边界框内的热力图defvisualize_bounding_box_heatmap(img,boxes,visualization):forboxinboxes:x1,y1,x2,y2int(box[0]),int(box[1]),int(box[2]),int(box[3])cropped_imgimg[y1:y2,x1:x2]cropped_tensorF.to_tensor(cropped_img).unsqueeze(0)cropped_visualizationget_gradcam_visualization(model,cropped_tensor,cam)img[y1:y2,x1:x2]cropped_visualizationreturnimg# 去除非红色区域defremove_non_red_regions(visualization):lower_rednp.array([200,0,0],dtypeuint8)upper_rednp.array([255,100,100],dtypeuint8)maskcv2.inRange(visualization,lower_red,upper_red)visualization[mask0][255,255,255]returnvisualization# 主程序if__name____main__:img_pathpath_to_your_image.jpgimg,img_tensorpreprocess_image(img_path)# 普通热力图visualizationget_gradcam_visualization(model,img_tensor,cam)plt.imshow(visualization)plt.show()# 仅标签框内的热力图boxesget_bounding_boxes(model,img_tensor)heatmap_with_boxesvisualize_bounding_box_heatmap(img,boxes,visualization)plt.imshow(heatmap_with_boxes)plt.show()# 去除红色以外区域的热力图heatmap_without_non_redremove_non_red_regions(visualization)plt.imshow(heatmap_without_non_red)plt.show()展示了如何在YOLOv5中实现三种不同的热力图可视化并且可以用于各种 YOLOv5模型以添加注意力机制和改进模型。仅供参考。