From f5c277124240165a7ddd7a75bbcaa0333d031d99 Mon Sep 17 00:00:00 2001 From: milesial Date: Wed, 6 Nov 2019 13:13:37 +0100 Subject: [PATCH] Fix predict script (mistake from last commit) Former-commit-id: 9fd386f5929945f3df1868930d392507fee16a55 --- predict.py | 10 ++++++++-- unet/unet_model.py | 5 ----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/predict.py b/predict.py index f8a2894..b26f655 100755 --- a/predict.py +++ b/predict.py @@ -6,6 +6,7 @@ import numpy as np import torch from PIL import Image from torchvision import transforms +import torch.nn.functional as F from unet import UNet from utils import plot_img_and_mask @@ -20,7 +21,6 @@ def predict_img(net, use_dense_crf=False): net.eval() img_height = full_img.size[1] - img_width = full_img.size[0] img = resize_and_crop(full_img, scale=scale_factor) img = normalize(img) @@ -32,7 +32,13 @@ def predict_img(net, with torch.no_grad(): output = net(X) - probs = output.squeeze(0) + + if net.n_classes > 1: + probs = F.softmax(output, dim=1) + else: + probs = torch.sigmoid(output) + + probs = probs.squeeze(0) tf = transforms.Compose( [ diff --git a/unet/unet_model.py b/unet/unet_model.py index 5b23b55..ee1f1bb 100644 --- a/unet/unet_model.py +++ b/unet/unet_model.py @@ -35,8 +35,3 @@ class UNet(nn.Module): x = self.up4(x, x1) logits = self.outc(x) return logits - - if self.n_classes > 1: - return F.softmax(x, dim=1) - else: - return torch.sigmoid(x)