From 54f0c1482427a5b3f2248b97be55878e742cbcb1 Mon Sep 17 00:00:00 2001
From: captin411 <captindave@gmail.com>
Date: Tue, 25 Oct 2022 16:14:13 -0700
Subject: [PATCH] download better face detection module dynamically

---
 modules/textual_inversion/autocrop.py   | 20 ++++++++++++++++++++
 modules/textual_inversion/preprocess.py | 13 +++++++++++--
 requirements.txt                        |  2 ++
 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