How to export a ContentDetectorFunction for use in Python?

Question

For text and images v13 introduces the ability to train detectors easy which result in ContentDetectorFunction's. To use these IRL we need to export to python's torch & onnx libraries.

I'm looking for guidance (or a tutorial) on how this can be done for both images (object detection) and text (entity detection). Here's an outline of what I've tried for the image case:

Running the first example from ref/TrainImageContentDetector:

It can take a while, so I uploaded the result for you:

imgdf = Import @ URLDownload["https://www.dropbox.com/s/c84ghzpkndwhyci/imgdf.mx?dl=1", "~/Downloads/imgdf.mx"];

The resulting ContentDetectorFunction object contains various keys:

But most important are "Net" and "Decoder". The first is a NetChain which exports fine, but it seems we can't Export FunctionLayer:

Export["~/Downloads/imgdf-net.onnx", imgdf[[1, "Net"]]]
Export["~/Downloads/imgdf-dec.onnx", NetGraph[imgdf[[1, "Decoder"]]]]

This is basically as far as I got, and I'm asking for help understanding the basics of pre-processing and post-processing that goes on with a ContentDetectorFunction during inference.

Notes:

• Most of the spelunking will be in the private context file19YOLO`.
• I'm very familiar with the onnx runtime and torch, so if a developer or someone can clarify the missing things on the mma side: encoding, non-max suppression, and decoding. Then I will write up the equivalent python code for those missing parts and post it here.

Answer 1

Let's start with the disclaimer that the internal code showed here is liable to change. I will try to keep this answer up-to-date but I cannot guarantee that =(

The image encoding for the YOLO method, which is the current default, is simply a letterboxing with gray padding. You can get the code from the NetEncoder via

NetExtract[imgdf[[1, "Net"]], {"Input", "Function"}]

Function[img,
    Catch[ImageData[
        ImageCrop[
            ColorConvert[
                ImageResize[
                    RemoveAlphaChannel[
                        Which[
                            ImageQ[img],
                                img
                            ,
                            FileExistsQ[img],
                                FastImageImport[img]
                            ,
                            True,
                                Throw[$Failed]
                        ]
                    ], {416}, Resampling -> "Linear"
                ], "RGB"
            ], {416, 416}, Padding -> 0.5`
        ], Interleaving -> False
    ]]
]

where you can replace FastImageImport with your preferred image library method.

The post processing is a little more convoluted. YOLO is an anchor based detector where the net predicts correction on some position priors (the cells in the final convolutions) and size/aspect ratio priors, which are typically called anchors. Therefore the output of the net is a grid where each element is a combination of

K * (4 [coordinates] + 
    1 [probability of object] + 
    nclasses [conditioned class probability])

K is the number of anchors used in the detector. If you need to reimplement it from scratch, the decoder doing the following:

• reshaping this array, from [K * (5 + nclasses), W, H] to [K, (5 + nclasses), W, H]
• then transposing to [(5 + nclasses), K, W, H]
• the probability is then sent to a logit function
• if there are more then one class, they are softmaxed
• the coordinates are in a center-size format (x, y, w, h)
  • sizes (part 3 and 4) are rescaled using the anchors via a$w * Exp[w]
  • centers (part 1 and 2) are considered offsets wrt the convolutional grid and their transformation is (pos$x * LogisticSigmoid[x]) / W

Once that is done you can work with a flat list of predictions. The final processing is

• filter detection with low probability (AcceptanceThreshold)
• filter more based on non-maximum suppression (nms), I put an implementation with code available in the function repository
• find the class with highest probability for each detection

If you have more classes you can imagine doing the nms globally or per class.

And that is more or less it. I hope it helps =)

I am going to edit this if there are questions in the comments.

About the measurements: I have some code to compute the mean average precision on a dataset and I would like to find a way to expose it in the next version.

Answer 2

Unclear if this answers the OP's question, and if so the process is rather manual, but here's some observations/steps to export the Decoder:

1. The OP correctly wrapped imgdf[[1, "Decoder"]] in a NetGraph to expand FunctionLayer out. However, there's a another nested FunctionLayer in there, hence the error message. The following flattens those out (manually):

ng=NetGraph[imgdf[[1,"Decoder"]]];
ngFlat=NetReplacePart[ng,"scaledBoxes"->NetGraph[NetExtract[ng,"scaledBoxes"]]];
Export["~/Downloads/imgdf-dec.onnx",ngFlat]

Unfortunately, the export line also fails with the error message:

Export::onnxbadparam4
Setting "Broadcasting" of ThreadingLayer can only be exported to ONNX when set to either 1, Automatic or None "

I suspect this is related to the Broadcasting/Listability inconsistency discussed here.

2. We can follow the error message's instruction and set Broadcasting to Automatic in the two ThreadingLayer instances it occurs

ngThreading=NetReplacePart[ngFlat,{"scaledBoxes","scaledCenters","Broadcasting"}->Automatic];
ngThreading=NetReplacePart[ngThreading,{"scaledBoxes","scaledSizes","Broadcasting"}->Automatic];
Export["~/Downloads/imgdf-dec.onnx",ngThreading]

This exports to ONNX fine, although I don't know if it still evaluates properly, i.e. if the swapping of the Broadcasting dimensions messes things up.