Spaces:
Configuration error
Configuration error
File size: 1,080 Bytes
0034848 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from __future__ import division
import numpy as np
import torch
import torchvision.transforms.functional as F
def img_to_tensor(im, normalize=None):
tensor = torch.from_numpy(np.moveaxis(im / (255.0 if im.dtype == np.uint8 else 1), -1, 0).astype(np.float32))
if normalize is not None:
return F.normalize(tensor, **normalize)
return tensor
def mask_to_tensor(mask, num_classes, sigmoid):
if num_classes > 1:
if not sigmoid:
# softmax
long_mask = np.zeros((mask.shape[:2]), dtype=np.int64)
if len(mask.shape) == 3:
for c in range(mask.shape[2]):
long_mask[mask[..., c] > 0] = c
else:
long_mask[mask > 127] = 1
long_mask[mask == 0] = 0
mask = long_mask
else:
mask = np.moveaxis(mask / (255.0 if mask.dtype == np.uint8 else 1), -1, 0).astype(np.float32)
else:
mask = np.expand_dims(mask / (255.0 if mask.dtype == np.uint8 else 1), 0).astype(np.float32)
return torch.from_numpy(mask)
|