Create model.py

model.py

@@ -0,0 +1,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)