简易方法采用CAM方法获取显著图和包围盒

╰半夏微凉° 2024-02-19 20:10 46阅读 0赞

# coding=utf-8
    # simple implementation of CAM in PyTorch for the networks such as ResNet, DenseNet, SqueezeNet, Inception
    import io
    import requests
    from PIL import Image
    from torchvision import models, transforms
    from torch.autograd import Variable
    from torch.nn import functional as F
    import numpy as np
    import cv2
    import json
    import scipy.io
    
    # input image #download picture online,or you can use your locate picture
    #LABELS_URL = 'https://s3.amazonaws.com/outcome-blog/imagenet/labels.json'
    #IMG_URL = 'http://media.mlive.com/news_impact/photo/9933031-large.jpg'
    
    # networks such as googlenet, resnet, densenet already use global average pooling at the end, so CAM could be used directly.
    model_id = 5  #选择一个预训练模型，也可添加其他模型，只需要知道模型名字和最后一个卷积层的名字
    if model_id == 1:
        net = models.squeezenet1_1(pretrained=True)
        finalconv_name = 'features' # this is the last conv layer of the network
    elif model_id == 2:
        net = models.resnet18(pretrained=True)
        finalconv_name = 'layer4'
    elif model_id == 3:
        net = models.densenet161(pretrained=True)
        finalconv_name = 'features'
    elif model_id == 4:
        net = models.resnet34(pretrained=True)
        finalconv_name =  'layer4'
    elif model_id == 5:
        net = models.resnet50(pretrained=True)
        finalconv_name =  'layer4'
    
    net.eval()
    
    # hook the feature extractor
    features_blobs = []
    def hook_feature(module, input, output):
        features_blobs.append(output.data.cpu().numpy())
    net._modules.get(finalconv_name).register_forward_hook(hook_feature)
    
    # get the softmax weight
    params = list(net.parameters())
    weight_softmax = np.squeeze(params[-2].data.numpy())
    
    def returnCAM(feature_conv, weight_softmax, class_idx):
        # generate the class activation maps upsample to 256x256
        size_upsample = (256, 256)
        bz, nc, h, w = feature_conv.shape
        output_cam = []
        for idx in class_idx:
            cam = weight_softmax[idx].dot(feature_conv.reshape((nc, h*w)))
            cam = cam.reshape(h, w)
            cam = cam - np.min(cam)
            cam_img = cam / np.max(cam)
            cam_img = np.uint8(255 * cam_img)
            output_cam.append(cv2.resize(cam_img, size_upsample))
        return output_cam
    
    #图片预处理
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
    preprocess = transforms.Compose([transforms.Resize((224,224)), transforms.ToTensor(), normalize])
    
    #response = requests.get(IMG_URL)
    #img_pil = Image.open(io.BytesIO(response.content))
    #使用本地图片
    test_img_name = 'Indigo_Bunting_0059_11596.jpg'  #此处是图片路径，因为我把图片放在当前目录下
    img_pil = Image.open(test_img_name)
    #img_pil.save('test.jpg')
    
    img_tensor = preprocess(img_pil)
    img_variable = Variable(img_tensor.unsqueeze(0))
    logit = net(img_variable)
    #print('logit:',logit)
    # download the imagenet category list
    #classes = {int(key):value for (key, value) in requests.get(LABELS_URL).json().items()}
    with open('labels.json','r') as json_class:
        classes = {int(key):value for (key, value) in json.load(json_class).items()}
    #print('classes:',classes) #classes: {0: 'tench, Tinca tinca', 1: 'goldfish, Carassius auratus',
    
    h_x = F.softmax(logit, dim=1).data.squeeze()
    #print('h_x: ',h_x)
    probs, idx = h_x.sort(0, True)
    probs = probs.numpy()
    idx = idx.numpy()
    #print('probs:',probs[:10])
    #print('idx',idx[:10])
    # output the prediction
    for i in range(0, 5):
        print('{:.3f} -> {}'.format(probs[i], classes[idx[i]]))
    
    # generate class activation mapping for the top1 prediction
    CAMs = returnCAM(features_blobs[0], weight_softmax, [idx[0]])#CAMs是一个包含一个array类型元素的list
    # cams shape: [array([[ 8,  8,  8, ..., 18, 18, 18],[ 8,  8,  8, ..., 18, 18, 18],[ 8,  8,  8, ..., 18, 18, 18], ...]
    
    print('cams[0] shape:',CAMs[0].shape)
    
    # render the CAM and output
    print('output CAM picture for the top1 prediction: %s'%classes[idx[0]])
    img = cv2.imread(test_img_name)
    img = cv2.resize(img,(256,256))
    height, width, _ = img.shape
    heatmap = cv2.applyColorMap(cv2.resize(CAMs[0],(width, height)), cv2.COLORMAP_JET)
    result = heatmap * 0.3 + img * 0.5
    cv2.imwrite('CAM_resnet50.jpg', result)
    
    #np.savetxt('cam.csv', CAMs[0], delimiter = ',') #保存CAMs[0]的值为.csv文件
    #获取阈值范围内的显著区域位置
    left,right,top,bottom = 0,0,0,0  #上线，下线，左线，右线
    cam_value = 60  # 可调整 Bounding box 的范围
    for left in range(255):
        if any(CAMs[0][:,left] >= cam_value):
            print('left:', left)
            break
    for right in range(255):
        if any(CAMs[0][:,255-right] >= cam_value):
            right = 255-right
            print('right:', right)
            break
    for top in range(255):
        if any(CAMs[0][top,:] >= cam_value):
            print('top:',top)
            break
    for bottom in range(255):
        if any(CAMs[0][255-bottom,:] >= cam_value):
            bottom = 255-bottom
            print('bottom:', bottom)
            break
    
    CAMs_np = np.array(CAMs[0])
    ROImap = CAMs_np[top:bottom,left:right]
    np.savetxt('ROImap.csv', ROImap, delimiter = ',') #保存ROImap的值为.csv文件
    
    cropped_Heatmap = heatmap[top:bottom, left:right]
    cv2.imwrite('croppedMap.jpg', cropped_Heatmap)
    #从原图中截取显著区域的图片并保存为 ROIimg.jpg
    img = cv2.resize(img,(256,256), interpolation=cv2.INTER_LINEAR)
    ROIimg = img[top:bottom, left:right]
    cv2.imwrite('ROIimg.jpg', ROIimg)

以上代码结合了CAM方法的官方实现以及自己的应用需求

cam论文:

[learning deep features for discriminative localization][]

[learning deep features for discriminative localization]: https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Zhou_Learning_Deep_Features_CVPR_2016_paper.pdf