Model Card

Model Description

• Model type: Computer vision - keypoint detection
• Fine-tuned from: keypointrcnn_resnet50_fpn

Get Started

1. Install PyTorch and Torchvision.
2. Instantiate the model and replace the prediction heads.
   Python

   1from torchvision.models.detection import keypointrcnn_resnet50_fpn, KeypointRCNN_ResNet50_FPN_Weights 2from torchvision.models.detection.faster_rcnn import FastRCNNPredictor 3from torchvision.models.detection.keypoint_rcnn import KeypointRCNNPredictor 4 5# Load a pre-trained Keypoint RCNN model 6weights = KeypointRCNN_ResNet50_FPN_Weights.DEFAULT 7model = keypointrcnn_resnet50_fpn(weights=weights) 8 9# Replace model's head 10num_classes = 2 11in_features = model.roi_heads.box_predictor.cls_score.in_features 12model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) 13 14num_keypoints = 14 15in_features = model.roi_heads.keypoint_predictor.kps_score_lowres.in_channels 16model.roi_heads.keypoint_predictor = KeypointRCNNPredictor(in_features, num_keypoints)
3. Download the model weight and load the state dict to the model.
   Python

   1from safetensors.torch import load_model 2load_model(model, "model.safetensors", device="cuda")
4. Refer to the Keypoint R-CNN doc for the model's usage.
5. Otherwise, refer to this script to export the model to ONNX and OpenVINO IR.