download better face detection module dynamically

3 files changed, 33 insertions, 2 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

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 1e4d4de8..e13b1894 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -7,6 +7,7 @@ import tqdm
 import time

 

 from modules import shared, images

+from modules.paths import models_path

 from modules.shared import opts, cmd_opts

 from modules.textual_inversion import autocrop

 if cmd_opts.deepdanbooru:

@@ -146,14 +147,22 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
                 save_pic(splitted, index, existing_caption=existing_caption)

             process_default_resize = False

 

-        if process_entropy_focus and img.height != img.width:

+        if process_focal_crop and img.height != img.width:

+

+            dnn_model_path = None

+            try:

+                dnn_model_path = autocrop.download_and_cache_models(os.path.join(models_path, "opencv"))

+            except Exception as e:

+                print("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", e)

+

             autocrop_settings = autocrop.Settings(

                 crop_width = width,

                 crop_height = height,

                 face_points_weight = process_focal_crop_face_weight,

                 entropy_points_weight = process_focal_crop_entropy_weight,

                 corner_points_weight = process_focal_crop_edges_weight,

-                annotate_image = process_focal_crop_debug

+                annotate_image = process_focal_crop_debug,

+                dnn_model_path = dnn_model_path,

             )

             for focal in autocrop.crop_image(img, autocrop_settings):

                 save_pic(focal, index, existing_caption=existing_caption)

diff --git a/requirements.txt b/requirements.txt
index da1969cf..75b37c4f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,6 +8,8 @@ gradio==3.5
 invisible-watermark

 numpy

 omegaconf

+opencv-python

+requests

 piexif

 Pillow

 pytorch_lightning