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])