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)

[ppwwyyxx]

The metadata for this dataset is missing the keypoint connection rules: https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html#metadata-for-datasets

should be something like:

KEYPOINT_CONNECTION_RULES = [
    # face
    ("left_ear", "left_eye", (102, 204, 255)),
    ("right_ear", "right_eye", (51, 153, 255)),
    ("left_eye", "nose", (102, 0, 204)),
    ("nose", "right_eye", (51, 102, 255)),
    # upper-body
    ("left_shoulder", "right_shoulder", (255, 128, 0)),
    ("left_shoulder", "left_elbow", (153, 255, 204)),
    ("right_shoulder", "right_elbow", (128, 229, 255)),
    ("left_elbow", "left_wrist", (153, 255, 153)),
    ("right_elbow", "right_wrist", (102, 255, 224)),
    # lower-body
    ("left_hip", "right_hip", (255, 102, 0)),
    ("left_hip", "left_knee", (255, 255, 77)),
    ("right_hip", "right_knee", (153, 255, 204)),
    ("left_knee", "left_ankle", (191, 255, 128)),
    ("right_knee", "right_ankle", (255, 195, 77)),
]

[jsperatan]

Thanks, that works!