download better face detection module dynamically

1 files changed, 20 insertions, 0 deletions

diff --git a/modules/textual_inversion/autocrop.py b/modules/textual_inversion/autocrop.py
index caaf18c8..01a92b12 100644
--- a/modules/textual_inversion/autocrop.py
+++ b/modules/textual_inversion/autocrop.py
@@ -1,4 +1,5 @@
 import cv2

+import requests

 import os

 from collections import defaultdict

 from math import log, sqrt

@@ -293,6 +294,25 @@ def is_square(w, h):
   return w == h

 

 

+def download_and_cache_models(dirname):

+  download_url = 'https://github.com/opencv/opencv_zoo/blob/91fb0290f50896f38a0ab1e558b74b16bc009428/models/face_detection_yunet/face_detection_yunet_2022mar.onnx?raw=true'

+  model_file_name = 'face_detection_yunet.onnx'

+

+  if not os.path.exists(dirname):

+    os.makedirs(dirname)

+

+  cache_file = os.path.join(dirname, model_file_name)

+  if not os.path.exists(cache_file):

+    print(f"downloading face detection model from '{download_url}' to '{cache_file}'")

+    response = requests.get(download_url)

+    with open(cache_file, "wb") as f:

+      f.write(response.content)

+

+  if os.path.exists(cache_file):

+    return cache_file

+  return None

+

+

 class PointOfInterest:

   def __init__(self, x, y, weight=1.0, size=10):

     self.x = x