Update app.py

app.py

@@ -14,6 +14,7 @@ from face_detection import RetinaFace
 from model import SixDRepNet
 import utils
 import cv2
+from PIL import Image
 
 snapshot_path = hf_hub_download(repo_id="osanseviero/6DRepNet_300W_LP_AFLW2000", filename="model.pth")
 
@@ -63,7 +64,9 @@ def predict(frame):
     p_pred_deg = euler[:, 0].cpu()
     y_pred_deg = euler[:, 1].cpu()
     r_pred_deg = euler[:, 2].cpu()
-    return utils.plot_pose_cube(frame, y_pred_deg, p_pred_deg, r_pred_deg, x_min + int(.5*(x_max-x_min)), y_min + int(.5*(y_max-y_min)), size = bbox_width)
+    arr = utils.plot_pose_cube(frame, y_pred_deg, p_pred_deg, r_pred_deg, x_min + int(.5*(x_max-x_min)), y_min + int(.5*(y_max-y_min)), size = bbox_width)
+    img = Image.fromarray(arr)
+    return img.resize(size=(200, 200))
  
 title = "6D Rotation Representation for Unconstrained Head Pose Estimation"
 description = "Gradio demo for 6DRepNet. To use it, simply click the camera picture. Read more at the links below."