Visualize Connected Keypoint Annotations

[jsperatan]

Hi I have successfully trained my own custom model and was using it to run inference on some images, The model appears to work pretty well and I have no issues with the display of the bounding box and the keypoint annotations itself. However I was expecting that the keypoints would be connected into a skeleton structure instead of just showing up as points. I have previously defined the connections when annotating using COCO annotator. Here is my code.

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints\keypoint_rcnn_R_50_FPN_1x.yaml"))
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # path to the model we just trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3   # set a custom testing threshold
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
predictor = DefaultPredictor(cfg)

register_coco_instances("infant", {}, "infant_annotation_train.json", "datasets\infant")

val_dir = r"validation/"
for file in os.listdir(val_dir):
    filename = os.fsdecode(file)
    filename = val_dir + filename
    im = cv2.imread(filename)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                metadata=MetadataCatalog.get("infant"),
                scale=0.5,
                instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models
                )
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    #out = v.draw_and_connect_keypoints(np.squeeze(outputs["instances"].pred_keypoints.to("cpu")))
    cv2.imshow("Inference", out.get_image()[:, :, ::-1])
    cv2.waitKey(0)

I have tried both the "draw_instance_predictions" and "draw_and_connect_keypoints" functions but both yielded similar results. (See image below for expectation vs reality)