Image segmentation and export to onnx


I’ve been playing with image segmentation and got it working pretty easily (yay fastai!!) but I’ve run into some stumbling block when trying to deploy that model. It works fine if I just package it with python and fastai but that is not what I would like like to do. My idea is to go mobile/browser and the only way I managed to do is by exporting to ONNX and that is where I’m stuck.

My ONNX model has last layer as:
%output1 : Float(1, 150, 224, 224) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1]](%337, %180, %181), scope: DynamicUnet/Sequential/Conv2d[0]

and I have no idea how to extract segmented mask from that.

I also noticed that learn.predict(img) returns:
[ImageSegment (1, 250, 250), torch.Size([1, 250, 250]), torch.Size([150, 250, 250]))
of which only the first item is segmented mask and the last corresponds basically to the output layer of the exported ONNX model.

I tried looking at the source of predict code but I don’t quite understand what is going on. :slight_smile:


1 Like