コード例 #1
0
def run_net_on_mrc_and_save(net_path,
                            protein_mrc_path,
                            prob_output_mrc_path=None,
                            binary_output_mrc_path=None):
    """
    loads a trained neural network, reads a cryo-EM map, runs the neural network on the cryo-EM map, applies a cutoff
    to the resulted map in order to round the values to either 0 or 1, and saves the original resulted map and
     he masked map to .mrc files.
    :param net_path: path of the trained neural network
    :param protein_mrc_path: path of the input cryo-EM map
    :param prob_output_mrc_path: path to save the .mrc prediction
    :param binary_output_mrc_path: path to save the .mrc binary mask of the prediction
    :return: the resulted helix prediction map
    """
    net = model.load_net(net_path)
    protein_data = torch.Tensor(
        dataset_manager.read_mrc(protein_mrc_path)).to('cuda')
    label_data = run_net_on_whole_protein(net, protein_data)
    actual_label_data_tensor = torch.Tensor(label_data).to('cuda')
    binary_label_data_tensor = dataset_manager.apply_cutoff(
        actual_label_data_tensor, 0.5)
    if prob_output_mrc_path != None:
        dataset_manager.save_mrc(label_data, prob_output_mrc_path)
    if binary_output_mrc_path != None:
        binary_label_data = binary_label_data_tensor.cpu().numpy()
        dataset_manager.save_mrc(binary_label_data, binary_output_mrc_path)
    # to calculate metrics uncomment the next line
    #metrics= validate_prediction(binary_label_data_tensor, dataset_manager.read_mrc(protein_mrc_path.split('.')[0] + '_helix.mrc')) # make it a comment later
    return label_data
コード例 #2
0
    '--threshold',
    type=float,
    default=0.2,
    help=
    'to suppress the corresponding CAM by the threshold of its maximum value')
parser.add_argument('--to_visualize',
                    type=bool,
                    default=True,
                    help='Save the CAMs as heatmap images for visualization')
parser.add_argument('--to_save',
                    type=bool,
                    default=False,
                    help='Save the CAMs as numpy array for further processing')
opt = parser.parse_args()

net = model.load_net(opt.model_name, opt.classes, opt.checkpoint_path).cuda()
finalconv_name = 'conv'

params = list(net.parameters())
# get weight only from the last layer(linear)
weight_softmax = params[-1].cpu().data


##========== Extract CAM for all classes with softmax normalization ==================#
def returnCAMs(feature_conv, weight_softmax, threshold=0.):
    bz, nc, h, w = feature_conv.shape
    features = feature_conv.reshape(nc, h * w).transpose(0, 1)  #(hxw, nc)
    CAMs = torch.mm(features, weight_softmax.transpose(0, 1))  #(hxw, k)
    CAMs = CAMs - torch.min(CAMs, dim=0)[0]
    CAMs = CAMs / torch.max(CAMs, dim=0)[0]
    # threshold the corresponding CAM by the set ratio of its maximum value
コード例 #3
0
ファイル: create_cam.py プロジェクト: srinidhigoud/CAM
from torch.nn import functional as F
import cv2
import numpy as np
import torchvision.transforms as transforms
import utils
import model

if not os.path.exists('./result'):
    os.mkdir('result/')

classes = ('airplance', 'bird', 'car', 'cat', 'deer', 'dog', 'horse', 'monkey',
           'ship', 'truck')
test_loader = utils.load_data_stl10(batch_size=1, test=True)
is_cuda = torch.cuda.is_available()
device = torch.device("cuda" if is_cuda else "cpu")
net = model.load_net().to(device)
finalconv_name = 'conv'

# hook
feature_blobs = []


def hook_feature(module, input, output):
    feature_blobs.append(output.cpu().data.numpy())


net._modules.get(finalconv_name).register_forward_hook(hook_feature)

params = list(net.parameters())
# get weight only from the last layer(linear)
weight_softmax = np.squeeze(params[-2].cpu().data.numpy())
コード例 #4
0
ファイル: create_cam.py プロジェクト: gchoi/CAM
import torch
import torch.nn as nn
from torch.autograd import Variable
import cv2
import numpy as np
import utils
import os
import torchvision.transforms as transforms
from torch.nn import functional as F

if os.path.exists('./result'):
    os.mkdir('result/')

test_loader = utils.load_data_STL10(batch_size=1, test=True)

net = model.load_net()
finalconv_name = 'conv'

# hook
feature_blobs = []


def hook_feature(module, input, output):
    feature_blobs.append(output.data.cpu().numpy())


net._modules.get(finalconv_name).register_forward_hook(hook_feature)

params = list(net.parameters())
# get only weight from last layer(linear)
weight_softmax = np.squeeze(params[-2].cpu().data.numpy())