# %%
if __name__ == '__main__':
# %%
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ckpt_path = EXP_DIR / 'mobile_seg/1607075632/checkpoints/last.ckpt'
loader = get_loader(
DataParams(
batch_size=1,
fold=0,
n_splits=5,
img_size=224,
seed=1,
))
model = load_trained_model(ckpt_path).to(device).eval()
# %%
with torch.no_grad():
inputs, labels = next(iter(loader))
outputs = model(inputs.to(device)).cpu()
inputs = inputs.squeeze()
labels = labels.squeeze()
outputs = outputs.squeeze()
inputs = (inputs * 255).numpy().transpose((1, 2, 0)).astype(np.uint8)
plt.subplot(131)
plt.imshow(inputs)
import coremltools as ct
import torch
from mobile_seg.const import TMP_DIR, EXP_DIR
from mobile_seg.modules.net import load_trained_model
from mobile_seg.modules.wrapper import Wrapper
# %%
IMG_SIZE = 224
TMP_ONNX = TMP_DIR / 'tmp.onnx'
CKPT_PATH = EXP_DIR / 'mobile_seg/1596704750/checkpoints/epoch=194.ckpt'
# %%
unet = load_trained_model(CKPT_PATH)
model = Wrapper(unet=unet).eval()
# %%
inputs = torch.randn(1, 3, IMG_SIZE, IMG_SIZE)
traced_model = torch.jit.trace(model, inputs)
# %%
# Convert to Core ML using the Unified Conversion API
# TODO 2020/08/06, It does not work, because coremltools does not support PyTorch 1.6 yet.
# See https://github.com/apple/coremltools/issues/827
model = ct.convert(
traced_model,
inputs=[ct.ImageType(name="input_1", shape=inputs.shape)],
)