coremltools conversion from pytorch to coreml

[vjdaemp]

@WillBrennan Do you have experience using coremltools to convert pytorch models to coreml?

I created a script to convert the "model_segmentation_skin_30.pth" model to coreml format, but have encountered issues. The coreml model is generated successfully, but the segmentation outputs are not working. I believe there may be an issue with the conversion script.

import urllib
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import torch
import torch.nn as nn
import torchvision
import json

from torchvision import transforms
from PIL import Image

from semantic_segmentation.models.fcn import FCNResNet101

from coremltools.converters.mil import register_torch_op
from coremltools.converters.mil.frontend.torch.ops import _get_inputs
from coremltools.converters.mil.mil import Builder as mb
import coremltools as ct

@register_torch_op
def type_as(context, node):
    inputs = _get_inputs(context, node)
    context.add(mb.cast(x=inputs[0], dtype='int32'), node.name)


labels = ['skin']

device = torch.device('cpu')

# Load the model 

model = FCNResNet101(categories=labels)
model.load_state_dict(torch.load("./pretrained/model_segmentation_skin_30.pth", map_location=device))
model.eval()

# Load a sample image (cat_dog.jpg)
input_image = Image.open("./test/test-img.jpg")
input_image.show()

preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])

input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)

with torch.no_grad():
    output = model(input_batch)['out'][0]
torch_predictions = output.argmax(0)



def display_segmentation(input_image, output_predictions):
    # Create a color pallette, selecting a color for each class
    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
    colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
    colors = (colors % 255).numpy().astype("uint8")

    # Plot the semantic segmentation predictions of 21 classes in each color
    r = Image.fromarray(
        output_predictions.byte().cpu().numpy()
    ).resize(input_image.size)
    r.putpalette(colors)

    # Overlay the segmentation mask on the original image
    alpha_image = input_image.copy()
    alpha_image.putalpha(255)
    r = r.convert("RGBA")
    r.putalpha(128)
    seg_image = Image.alpha_composite(alpha_image, r)
    # display(seg_image) -- doesn't work
    seg_image.show()

display_segmentation(input_image, torch_predictions)

# Wrap the Model to Allow Tracing*
class WrappedFCNResNet(nn.Module):
    
    def __init__(self):
        super(WrappedFCNResNet, self).__init__()
        self.model = FCNResNet101(categories=labels)
        self.model.load_state_dict(torch.load("./pretrained/model_segmentation_skin_30.pth", map_location=device))
        self.model.eval()
    def forward(self, x):
        res = self.model(x)
        x = res["out"]
        return x
        
# Trace the Wrapped Model
traceable_model = WrappedFCNResNet().eval()
trace = torch.jit.trace(traceable_model, input_batch)

# Convert the model
mlmodel = ct.convert(
    trace,
    inputs=[ct.TensorType(name="input", shape=input_batch.shape)],
)

# Save the model without new metadata
mlmodel.save("SkinSegmentation_no_metadata.mlmodel")

# Load the saved model
mlmodel = ct.models.MLModel("SkinSegmentation_no_metadata.mlmodel")

# Add new metadata for preview in Xcode
labels_json = {"labels": ["skin"]}

mlmodel.user_defined_metadata["com.apple.coreml.model.preview.type"] = "SkinSegmentation"
mlmodel.user_defined_metadata['com.apple.coreml.model.preview.params'] = json.dumps(labels_json)

mlmodel.save("SkinSegmentation_plus_metadata.mlmodel")

Does anything look off to you in above script?