torch.squeeze(model(img), dim=0).detach().cpu()
这是一个 PyTorch 模型的推理代码,其中 model
是模型对象,img
是输入的图像数据。torch.squeeze
函数可以将维度为 1 的维度压缩掉,dim
参数指定要压缩的维度。detach
函数可以将计算图中的某个节点分离出来,使得该节点不再参与梯度计算。cpu
函数可以将张量从 GPU 上移动到 CPU 上。最终返回的是一个 CPU 上的张量,表示模型对输入图像的推理结果。
import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms as transforms from PIL import Image # 加载图像 img = Image.open('2.jpg') # 对图像进行预处理,将其转换为模型所需的输入格式 transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) img_tensor = transform(img).unsqueeze(0) # 初始化模型并对图像进行特征提取 model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', pretrained=True) features = model.conv1(img_tensor) features = model.bn1(features) features = model.relu(features) features = model.maxpool(features) features = model.layer1(features) features = model.layer2(features) features = model.layer3(features) features = model.layer4(features) # 将特征图还原回原始图像大小 upsample = nn.Upsample(scale_factor=32, mode='bilinear', align_corners=True) upsampled_features = upsample(features) # 显示原始图像和还原后的特征图 img.show() tensor_to_image = transforms.ToPILImage() upsampled_image = tensor_to_image(upsampled_features.squeeze(0).detach().cpu()) upsampled_image.show(),上述代码出现问题:ValueError: pic should not have > 4 channels. Got 2048 channels.
这个问题可能是因为在对特征图进行还原时,upsampled_features
的通道数为2048,而ToPILImage()
只能处理通道数为1、3或4的图像。解决方法是在还原特征图之前,将其通道数降为3,例如可以使用torch.mean
将特征图的通道数降为1,然后再使用torch.cat
将其复制为3个通道。具体代码如下:
# 将特征图还原回原始图像大小
upsample = nn.Upsample(scale_factor=32, mode='bilinear', align_corners=True)
upsampled_features = upsample(features)
# 将通道数降为1
feature_map = torch.mean(upsampled_features, dim=1, keepdim=True)
# 将通道数复制为3
feature_map = torch.cat([feature_map]*3, dim=1)
# 显示原始图像和还原后的特征图
img.show()
tensor_to_image = transforms.ToPILImage()
upsampled_image = tensor_to_image(feature_map.squeeze(0).detach().cpu())
upsampled_image.show()
这样就可以将特征图还原并显示出来了。
def predict_one_img(model, img_dataset, args): # def predict_one_img(model, img_dataset, data, target, args): dataloader = DataLoader(dataset=img_dataset, batch_size=1, num_workers=0, shuffle=False) model.eval() test_dice = DiceAverage(args.n_labels) # target = to_one_hot_3d(label, args.n_labels) with torch.no_grad(): for data ,target in tqdm(dataloader, total=len(dataloader)): # data = data.to(device) data, target = data.float(), target.long() target = to_one_hot_3d(target, args.n_labels) data, target = data.to(device), target.to(device) # print(data.shape) # print(target.shape) output = model(data) # output = nn.functional.interpolate(output, scale_factor=(1//args.slice_down_scale,1//args.xy_down_scale,1//args.xy_down_scale), mode='trilinear', align_corners=False) # 空间分辨率恢复到原始size img_dataset.update_result(output.detach().cpu()) pred = img_dataset.recompone_result() pred = torch.argmax(pred, dim=1) pred_img = to_one_hot_3d(pred, args.n_labels) pred_img=pred_img.to(device) test_dice.update(pred_img, target) test_dice = OrderedDict({'Dice_liver': test_dice.avg[1]}) if args.n_labels == 3: test_dice.update({'Dice_tumor': test_dice.avg[2]}) pred = np.asarray(pred.numpy(), dtype='uint8') if args.postprocess: pass # TO DO pred = sitk.GetImageFromArray(np.squeeze(pred, axis=0)) return test_dice, pred
这是一个用于预测单张图像的函数,接收一个模型、一个图像数据集和一些参数作为输入。函数首先将图像数据集加载到一个Dataloader中,然后将模型设置为评估模式并初始化一个DiceAverage对象,该对象用于计算评估指标(这里是Dice系数)。在没有梯度的情况下,对于每个数据和目标对,函数将数据和目标转换为模型所需的格式(浮点数和one-hot编码),并将它们传递给模型进行预测。预测结果被转换回原来的空间分辨率,然后被添加到图像数据集中。最后,预测结果被重新组合成一个numpy数组,并被转换为一个SimpleITK图像对象并返回。如果需要,可以在预测之后对预测结果进行后处理。函数返回测试指标和预测结果。