Hi, I generated an onnx file for keypoint and then further generating the converted onnx for tensorrt engine.
Thing is that in the output, instead of keypoint x, y points, I get the keypoint heatmap values.
I will attach the onnx and the converted onnx.
Onnx:
Converted Onnx:
If you see the onnx file, at node ConstantOfShape_2057, I can see the xy_preds .
But on the converted onnx at the last node (gather), the output dimension is 100x17x56x56 which I think are heatmap values since I believe xy_preds should be 100x17x3x3.

I modified the converted onnx python code by doing something like this. (took reference from TensorRT’s github)
mask_pooler_output = self.ROIAlign(nms_outputs[1], p2, p3, p4, p5, self.second_ROIAlign_pooled_size, \
self.second_ROIAlign_sampling_ratio, self.second_ROIAlign_type, self.second_NMS_max_proposals, 'keypoint_pooler')
# Reshape mask pooler output.
mask_pooler_shape = np.asarray([self.second_NMS_max_proposals*self.batch_size, self.fpn_out_channels, self.second_ROIAlign_pooled_size, self.second_ROIAlign_pooled_size], dtype=np.int64)
mask_pooler_reshape_node = self.graph.op_with_const("Reshape", "keypoint_pooler/reshape", mask_pooler_output, mask_pooler_shape)
# Get first Conv op in mask head and connect ROIAlign's squeezed output to it.
mask_head_conv = self.graph.find_node_by_op_name("Conv", "/roi_heads/keypoint_head/conv_fcn1/Conv")
mask_head_conv.inputs[0] = mask_pooler_reshape_node[0]
# Reshape node that is preparing 2nd NMS class outputs for Add node that comes next.
classes_reshape_shape = np.asarray([self.second_NMS_max_proposals * self.batch_size], dtype=np.int64)
classes_reshape_node = self.graph.op_with_const("Reshape", "box_outputs/reshape_classes", nms_outputs[3], classes_reshape_shape)
# This loop will generate an array used in Add node, which eventually will help Gather node to pick the single
# class of interest per bounding box, instead of creating 80 masks for every single bounding box.
add_array = []
for i in range(self.second_NMS_max_proposals * self.batch_size):
if i == 0:
start_pos = 0
else:
start_pos = i * self.num_classes
add_array.append(start_pos)
# This Add node is one of the Gather node inputs, Gather node performs gather on 0th axis of data tensor
# and requires indices that set tensors to be withing bounds, this Add node provides the bounds for Gather.
add_array = np.asarray(add_array, dtype=np.int32)
classes_add_node = self.graph.op_with_const("Add", "box_outputs/add", classes_reshape_node[0], add_array)
# Get the last Conv op in mask head and reshape it to correctly gather class of interest's masks.
last_resize = self.graph.find_node_by_op_name("Resize", "/roi_heads/keypoint_head/Resize")
# Gather node that selects only masks belonging to detected class, 79 other masks are discarded.
final_gather = self.graph.gather("/keypoint_head/gathering", last_resize.outputs[0], classes_add_node[0])
final_gather[0].dtype = np.float32
return nms_outputs, final_gather[0`
Could you please help me with the converted onnx? Maybe I am doing something wrong?
EDIT:
From what I understand the problem could be is that, xy_pred stuff happens after the resize node but in the code, after getting the last resize, I am just gathering it.
Somehow I want to insert the node ConstantOfShape_2057 but whenever I try to insert it in the code, it throws error.
# Get the last Conv op in mask head and reshape it to correctly gather class of interest's masks.
last_resize = self.graph.find_node_by_op_name("Resize", "/roi_heads/keypoint_head/Resize")
Thanks.