from blendbatchnorm import fuse_bn_recursively

import onnx

from onnx2keras import onnx_to_keras

import torch

import numpy as np

from meshnet import MeshNet

import tensorflowjs as tfjs

from fixmodeljson import fixjson_file

def preprocess_image(img):

    """Unit interval preprocessing"""

    img = (img - img.min()) / (img.max() - img.min())

    return img

volume_shape = [256, 256, 256]

subvolume_shape = [38, 38, 38]

n_subvolumes = 1024

n_classes = 3

atlas_classes = 104

scube = 64

model_path = '../meshnet_gmwm_dropout_train.30_full.pth'

#'meshnet_gmwm_dropout_train.30_full.pth'#'meshnet_gmwm_train.30_full.pth'

device_name = "cuda:0" if torch.cuda.is_available() else "cpu"

device = torch.device(device_name)

meshnet_model = MeshNet(n_channels=1, n_classes=n_classes, large=False)

meshnet_model.load_state_dict(torch.load(model_path, map_location=device)['model_state_dict'])

meshnet_model.to(device)

mnm = fuse_bn_recursively(meshnet_model)

mnm.model.eval();

x = torch.randn(1, 1, scube, scube, scube, requires_grad=True)

torch.onnx.export(mnm, x.to(device), '/tmp/mnm_model_large.onnx', export_params=True, opset_version=13, do_constant_folding=True, input_names = ['input'], output_names = ['output'],dynamic_axes={'input' : {0 : 'batch_size'},'output' : {0 : 'batch_size'}})

onnx_model = onnx.load('/tmp/mnm_model_large.onnx')

k_model = onnx_to_keras(onnx_model, ['input'])

tfjs.converters.save_keras_model(k_model, '/tmp/mnm_gmwm_dropout256')

fixjson_file('/tmp/mnm_gmwm_dropout256/model.json', scube=scube)