def predict(net, full_img, device, input_size, mask_way='warp'):
'''
:mask_type: Sets the way to obtain the mask. Сan take 'warp' or 'segm'
'''
# Preprocess input image:
img = BasicDataset.preprocess_img(full_img, input_size)
img = img.unsqueeze(0)
img = img.to(device=device, dtype=torch.float32)
net.eval()
# Predict:
with torch.no_grad():
logits, rec_mask, theta = net.predict(
img, warp=True if mask_way == 'warp' else False)
if mask_way == 'warp':
mask = rec_mask * net.n_classes
mask = mask.type(torch.IntTensor).cpu().numpy().astype(np.uint8)
elif mask_way == 'segm':
mask = preds_to_masks(logits, net.n_classes)
else:
raise NotImplementedError
return mask, theta
def predict_img(net, full_img, device, input_size):
net.eval()
img = BasicDataset.preprocess_img(full_img, input_size)
img = img.unsqueeze(0)
img = img.to(device=device, dtype=torch.float32)
with torch.no_grad():
preds = net(img)
masks = preds_to_masks(preds, net.n_classes) # GPU tensor -> CPU numpy
return masks
def predict_img(net, full_img, device, input_size):
# Preprocess input image:
img = BasicDataset.preprocess_img(full_img, input_size)
img = img.unsqueeze(0)
img = img.to(device=device, dtype=torch.float32)
net.eval()
# Predict:
with torch.no_grad():
mask_pred, mask_proj = net(img)
# Tensors to ndarrays:
mask = preds_to_masks(mask_pred, net.n_classes)
proj = mask_proj * net.n_classes
proj = proj.type(torch.IntTensor).cpu().numpy().astype(np.uint8)
return mask, proj
