Update

app.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import functools
 import sys
 from typing import Callable
 
@@ -66,25 +65,6 @@ def crop_face(image: np.ndarray, box: tuple[int, int, int, int]) -> np.ndarray:
     return image
 
 
-@spaces.GPU
-@torch.inference_mode()
-def predict(image: np.ndarray, transform: Callable, model: nn.Module, device: torch.device) -> np.ndarray:
-    indices = torch.arange(66).float().to(device)
-
-    image = PIL.Image.fromarray(image)
-    data = transform(image)
-    data = data.to(device)
-
-    # the output of the model is a tuple of 3 tensors (yaw, pitch, roll)
-    # the shape of each tensor is (1, 66)
-    out = model(data[None, ...])
-    out = torch.stack(out, dim=1)  # shape: (1, 3, 66)
-    out = F.softmax(out, dim=2)
-    out = (out * indices).sum(dim=2) * 3 - 99
-    out = out.cpu().numpy()[0]
-    return out
-
-
 def draw_axis(image: np.ndarray, pose: np.ndarray, origin: np.ndarray, length: int) -> None:
     # (yaw, pitch, roll) -> (roll, yaw, pitch)
     pose = pose[[2, 0, 1]]
@@ -99,19 +79,33 @@ def draw_axis(image: np.ndarray, pose: np.ndarray, origin: np.ndarray, length: i
     cv2.line(image, tuple(origin), tuple(pts[2]), (255, 0, 0), 2)
 
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+face_detector = RetinaFacePredictor(threshold=0.8, device="cpu", model=RetinaFacePredictor.get_model("mobilenet0.25"))
+face_detector.device = device
+face_detector.net.to(device)
+
+model_names = [
+    "hopenet_alpha1",
+    "hopenet_alpha2",
+    "hopenet_robust_alpha1",
+]
+models = {name: load_model(name, device) for name in model_names}
+transform = create_transform()
+
+
+@spaces.GPU
+@torch.inference_mode()
 def run(
     image: np.ndarray,
     model_name: str,
-    face_detector: RetinaFacePredictor,
-    models: dict[str, nn.Module],
-    transform: Callable,
-    device: torch.device,
 ) -> np.ndarray:
     model = models[model_name]
 
     # RGB -> BGR
     det_faces = face_detector(image[:, :, ::-1], rgb=False)
 
+    indices = torch.arange(66).float().to(device)
+
     res = image[:, :, ::-1].copy()
     for det_face in det_faces:
         box = np.round(det_face[:4]).astype(int)
@@ -119,8 +113,17 @@ def run(
         # RGB
         face_image = crop_face(image, box.tolist())
 
-        # (yaw, pitch, roll)
-        angles = predict(face_image, transform, model, device)
+        face_image = PIL.Image.fromarray(face_image)
+        data = transform(face_image)
+        data = data.to(device)
+
+        # the output of the model is a tuple of 3 tensors (yaw, pitch, roll)
+        # the shape of each tensor is (1, 66)
+        out = model(data[None, ...])
+        out = torch.stack(out, dim=1)  # shape: (1, 3, 66)
+        out = F.softmax(out, dim=2)
+        out = (out * indices).sum(dim=2) * 3 - 99
+        angles = out.cpu().numpy()[0]
 
         center = (box[:2] + box[2:]) // 2
         length = (box[3] - box[1]) // 2
@@ -129,21 +132,6 @@ def run(
     return res[:, :, ::-1]
 
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-face_detector = RetinaFacePredictor(threshold=0.8, device="cpu", model=RetinaFacePredictor.get_model("mobilenet0.25"))
-face_detector.device = device
-face_detector.net.to(device)
-
-model_names = [
-    "hopenet_alpha1",
-    "hopenet_alpha2",
-    "hopenet_robust_alpha1",
-]
-models = {name: load_model(name, device) for name in model_names}
-transform = create_transform()
-
-fn = functools.partial(run, face_detector=face_detector, models=models, transform=transform, device=device)
-
 examples = [["images/pexels-ksenia-chernaya-8535230.jpg", "hopenet_alpha1"]]
 
 with gr.Blocks(css="style.css") as demo:
@@ -159,10 +147,10 @@ with gr.Blocks(css="style.css") as demo:
         examples=examples,
         inputs=[image, model_name],
         outputs=result,
-        fn=fn,
+        fn=run,
     )
     run_button.click(
-        fn=fn,
+        fn=run,
         inputs=[image, model_name],
         outputs=result,
         api_name="run",