[feature request] classification by axis support in classification extension

[dumbPy]

Is your feature request related to a problem? Please describe.

For models like object detections, there might be more than one classification operation required. eg, for a coco object detection model, it might output 3 bounding boxes where each bounding box has 80 probabilities corresponding to each of 80 labels. The output shape would be 3x80 (assuming bbox coordinates are a separate tensor) in this case. since there can be any number of objects, the output shape in config.pbtxt would be [-1, 80].

Accordingly to https://github.com/triton-inference-server/server/issues/3804#issuecomment-1022477772
Currently, classification extensions would flatten the (3,80) tensor into a 240 dim tensor and expects the labelmap to have 240+ labels (https://github.com/triton-inference-server/server/issues/3804#issuecomment-1013403250).
We need a way to apply classification across one axis of the output, in order to support object detection and segmentation labels.

Describe the solution you'd like

This can be achieved by introducing classification_axis parameter to output tensors in config.pbtxt across which the classification is to be applied. so axis=1 would mean apply classification across axis 1. not having this parameter set may result in current default behaviour of flattening the tensor and then classifying.

This would also solve https://github.com/triton-inference-server/server/issues/3804#issuecomment-1013170745 by applying classification with classification_axis=2 to perform classification across the 2 outputs for each of the 28x28 pixel the model outputs.

Also, this behaviour should change how class_count works. setting class_count=2 should return top 2 classes from the classification_axis. So in the above example of coco detection with [3,80] output probabilities, it should return [3,2] tensor with top-2 classification for each of the 3 boxes.

eg.

# output tensor of shape [3,80]
[[0.01,0.625,0.12,0.003,.....,],
 [0.002,0.03,...,0.4,0.88,...],
 [0.003,0.02,0.32,....,0.8....]]

should return, with class_count=2 and classification_axis=1

[[b'0.625:1:bicycle',b'0.12:2:car'],
 [b'0.88:16:dog', b'0.4:15:cat'],
 [b'0.8:6:truck', b'0.32:2:car']

Describe alternatives you've considered

There are no alternatives for using clasification extension with object detection. #3484 might help with assigning the labels on client side, but this feature request would eliminate the need for it altogether.

Additional context None.

[GuanLuo]

I think the feature request in this issue is similar to the one in #3644, we do have a ticket to track that issue and it would be good to know if that may fit your need.

[dzier]

I have added this issue to our internal ticket as a second source.

[dumbPy]

Yes #3644 is asking for something similar

[dumbPy]

Is this part of any roadmap or has any ETA?

The onnx/tensorrt/pytorch/tensorflow backend work great with all sort of models but since the labels are not supported, I have to wrap all models with python for this, and that is an unnecessary complication. Also, the ensemble model doesn't work with kserve/modelmesh so have to resort to python model with custom env and load onnxruntime inside it just because there's no support for labels in onnx/etc backend.

Would be a great help.