I’m trying to convert a Detectron2 model to onnx. When I try to run the onnx model afterwards I get the following error:

Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from /dbfs/FileStore/ECONSTRUCT/FloorplanAnalyser/mask_rcnn_R_101_FPN_3x/bitmask/model_simp.onnx failed:Node (If_3504) Op (If) [TypeInferenceError] Mismatched tensor element type: inferred=bool declared=uint8

The problem is related to the masks the model returns. I believe they are uint8 in pytorch and onnx expects boolean (or the other way around). When I exclude the masks the onnx model runs without error.

Here’s the code:

import torch
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.export import TracingAdapter
import onnx
import cv2
import onnxruntime as ort
import numpy as np

cfg = get_cfg()
model_file_name = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
cfg.merge_from_file(model_zoo.get_config_file(model_file_name))
cfg.MODEL.DEVICE = 'cpu'
cfg.MODEL.WEIGHTS = 'model_final.pth'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
cfg.TEST.EVAL_PERIOD = 200
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.INPUT.MASK_FORMAT= 'bitmask'
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.MAX_ITER = 2000
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3

class TorchModel(torch.nn.Module):
    def __init__(self, cfg) -> None:
        super().__init__()
        self.model = build_model(cfg) # Build Model
        _ = DetectionCheckpointer(self.model).load(cfg.MODEL.WEIGHTS)  # Load weights
        self.model = self.model.eval() # In evaluation mode
    
    def forward(self, INPUT):
        if isinstance(INPUT, (np.ndarray, torch.Tensor)): # it supports just 1 image
            INPUT = [{"image":INPUT}]

        with torch.no_grad():
            outputs = self.model(INPUT)[0]['instances']
        
        boxes, labels, scores, masks = outputs.pred_boxes.tensor, outputs.pred_classes, outputs.scores.detach(), outputs.pred_masks.to(dtype=torch.uint8)

        return boxes, labels, scores, masks
model = TorchModel(cfg)
model = model.eval()

metadata = MetadataCatalog.get(VALID_DATA_SET_NAME)
dataset_valid = DatasetCatalog.get(VALID_DATA_SET_NAME)
im = cv2.imread(dataset_valid[7]['file_name'])
im_torch = torch.as_tensor(im.astype("float32").transpose(2, 0, 1))
inputs = [{"image": im_torch}]

traceable_model = TracingAdapter(model, inputs, None)

torch.onnx.export(traceable_model, (im_torch,),
                  "model.onnx", 
                  opset_version=17,
                  input_names = ['image'],
                  output_names = ['boxes', 'labels', 'scores', 'mask'],
                  dynamic_axes={'image' : {1 : 'height', 2: 'width'},
                                'boxes' : {0 : 'num_boxes'},
                                'labels' : {0 : 'num_boxes'},
                                'scores' : {0 : 'num_boxes'},
                                'masks' : {0 : 'num_boxes', 1 : 'mask_height', 2 : 'mask_width'}
                                })

session = ort.InferenceSession("model.onnx")
im_test = cv2.imread('demo_image.png')
im_test_torch = torch.as_tensor(im_test.astype("float32").transpose(2, 0, 1))

im_test_numpy = im_test_torch.detach().numpy()

outputs = session.run(None, {"image": im_test_numpy})

As you can see I already tried to cast the masks to uint8, without success. Does anyone know how to solve this problem?

Below are some of the last nodes of the onnx model graph:

Node: Squeeze_3500, OpType: Squeeze, Inputs: ['onnx::Squeeze_3566', 'onnx::Squeeze_3567'], Outputs: ['N']
Node: Constant_3501, OpType: Constant, Inputs: [], Outputs: ['onnx::Equal_3569']
Node: Equal_3502, OpType: Equal, Inputs: ['N', 'onnx::Equal_3569'], Outputs: ['onnx::Cast_3570']
Node: Cast_3503, OpType: Cast, Inputs: ['onnx::Cast_3570'], Outputs: ['onnx::If_3571']
Node: If_3504, OpType: If, Inputs: ['onnx::If_3571'], Outputs: ['bitmasks']
Node: /model/model/Cast_11, OpType: Cast, Inputs: ['bitmasks'], Outputs: ['/model/model/Cast_11_output_0']
Node: /model/Cast, OpType: Cast, Inputs: ['/model/model/Cast_11_output_0'], Outputs: ['mask']