39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
|
from PIL import Image
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torch
|
||
|
|
import sys
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
imgSize = (1,2,640,959)
|
||
|
|
|
||
|
|
def mask_to_image(mask: np.ndarray, mask_values):
|
||
|
|
if isinstance(mask_values[0], list):
|
||
|
|
out = np.zeros((mask.shape[-2], mask.shape[-1], len(mask_values[0])), dtype=np.uint8)
|
||
|
|
elif mask_values == [0, 1]:
|
||
|
|
out = np.zeros((mask.shape[-2], mask.shape[-1]), dtype=bool)
|
||
|
|
else:
|
||
|
|
out = np.zeros((mask.shape[-2], mask.shape[-1]), dtype=np.uint8)
|
||
|
|
|
||
|
|
if mask.ndim == 3:
|
||
|
|
mask = np.argmax(mask, axis=0)
|
||
|
|
|
||
|
|
for i, v in enumerate(mask_values):
|
||
|
|
out[mask == i] = v
|
||
|
|
|
||
|
|
return Image.fromarray(out)
|
||
|
|
|
||
|
|
if __name__=='__main__':
|
||
|
|
|
||
|
|
state_dict = torch.load('MODEL.pth', map_location=torch.device('cpu') )
|
||
|
|
mask_values = state_dict.pop('mask_values', [0, 1])
|
||
|
|
|
||
|
|
img = np.fromfile(sys.argv[1], dtype=np.float32)
|
||
|
|
img = torch.from_numpy(img).reshape(imgSize)
|
||
|
|
|
||
|
|
img = F.interpolate(img, (640*2, 959*2), mode='bilinear')
|
||
|
|
mask = img.argmax(dim=1)
|
||
|
|
|
||
|
|
mask = mask[0].long().squeeze().numpy()
|
||
|
|
result = mask_to_image(mask, mask_values)
|
||
|
|
result.save(sys.argv[2])
|