Spaces:
Sleeping
Sleeping
File size: 2,897 Bytes
a8badcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | 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) |