Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| from torchvision.transforms import functional as F | |
| from torchvision.models.detection import ( | |
| fasterrcnn_resnet50_fpn_v2, | |
| keypointrcnn_resnet50_fpn, | |
| ) | |
| from torchvision.models.detection.faster_rcnn import FastRCNNPredictor | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # ----------------------------- | |
| # BUILD MODELS | |
| # ----------------------------- | |
| def build_face_detector(num_classes=2): | |
| model = fasterrcnn_resnet50_fpn_v2(weights=None) | |
| in_features = model.roi_heads.box_predictor.cls_score.in_features | |
| model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) | |
| return model | |
| def build_landmark_model(num_classes=2, num_keypoints=13): | |
| model = keypointrcnn_resnet50_fpn( | |
| weights=None, | |
| weights_backbone=None, | |
| num_classes=num_classes, | |
| num_keypoints=num_keypoints, | |
| ) | |
| return model | |
| # ----------------------------- | |
| # LOAD WEIGHTS | |
| # ----------------------------- | |
| def load_models(face_path, landmark_path, num_keypoints): | |
| face_model = build_face_detector() | |
| landmark_model = build_landmark_model(num_keypoints=num_keypoints) | |
| face_model.load_state_dict(torch.load(face_path, map_location=DEVICE)) | |
| landmark_model.load_state_dict(torch.load(landmark_path, map_location=DEVICE)) | |
| face_model.to(DEVICE).eval() | |
| landmark_model.to(DEVICE).eval() | |
| return face_model, landmark_model | |
| # ----------------------------- | |
| # CASCADE INFERENCE | |
| # ----------------------------- | |
| def run_inference(image, face_model, landmark_model, | |
| face_score_thr=0.5, | |
| kpt_score_thr=0.2): | |
| image_pil = Image.fromarray(image).convert("RGB") | |
| img_rgb = np.array(image_pil) | |
| img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) | |
| with torch.no_grad(): | |
| det_out = face_model([F.to_tensor(image_pil).to(DEVICE)])[0] | |
| boxes = det_out["boxes"].cpu().numpy() | |
| scores = det_out["scores"].cpu().numpy() | |
| keep = np.where(scores >= face_score_thr)[0] | |
| for i in keep: | |
| x1, y1, x2, y2 = boxes[i].astype(int) | |
| cv2.rectangle(img_bgr, (x1, y1), (x2, y2), (255, 180, 0), 2) | |
| crop = img_rgb[y1:y2, x1:x2] | |
| with torch.no_grad(): | |
| kp_out = landmark_model([F.to_tensor(crop).to(DEVICE)])[0] | |
| if len(kp_out["scores"]) == 0: | |
| continue | |
| best = torch.argmax(kp_out["scores"]).item() | |
| if kp_out["scores"][best] < kpt_score_thr: | |
| continue | |
| keypoints = kp_out["keypoints"][best].cpu().numpy() | |
| for kx, ky, kv in keypoints: | |
| if kv > 0: | |
| cv2.circle(img_bgr, | |
| (int(kx + x1), int(ky + y1)), | |
| 2, | |
| (0, 255, 0), | |
| -1) | |
| return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) |