| from detectron2.config import get_cfg |
| from detectron2.engine import DefaultPredictor |
| import cv2 |
| from detectron2 import model_zoo |
|
|
| def build_config(weights_path): |
| cfg = get_cfg() |
| |
| cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")) |
| cfg.MODEL.WEIGHTS = weights_path |
| cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 16 |
| cfg.MODEL.DEVICE = 'cpu' |
| cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 |
| cfg.MODEL.RETINANET.NUM_CLASSES = 1 |
|
|
| return cfg |
|
|
|
|
| def run_detection(image_path, weights_path): |
| cfg = build_config(weights_path) |
| |
| predictor = DefaultPredictor(cfg) |
| image = cv2.imread(image_path) |
| |
| outputs = predictor(image) |
| |
| instances = outputs["instances"] |
| |
| if len(instances) == 0: |
| return 'no keypoints' |
| else: |
| highest_confidence_index = instances.scores.argmax() |
|
|
| |
| highest_confidence_keypoints = instances.pred_keypoints[highest_confidence_index].cpu().numpy() |
|
|
| return highest_confidence_keypoints |