Commit 54f0c148 authored by captin411's avatar captin411

download better face detection module dynamically

parent db8ed5fe
import cv2 import cv2
import requests
import os import os
from collections import defaultdict from collections import defaultdict
from math import log, sqrt from math import log, sqrt
...@@ -293,6 +294,25 @@ def is_square(w, h): ...@@ -293,6 +294,25 @@ def is_square(w, h):
return 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: class PointOfInterest:
def __init__(self, x, y, weight=1.0, size=10): def __init__(self, x, y, weight=1.0, size=10):
self.x = x self.x = x
......
...@@ -7,6 +7,7 @@ import tqdm ...@@ -7,6 +7,7 @@ import tqdm
import time import time
from modules import shared, images from modules import shared, images
from modules.paths import models_path
from modules.shared import opts, cmd_opts from modules.shared import opts, cmd_opts
from modules.textual_inversion import autocrop from modules.textual_inversion import autocrop
if cmd_opts.deepdanbooru: if cmd_opts.deepdanbooru:
...@@ -146,14 +147,22 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre ...@@ -146,14 +147,22 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
save_pic(splitted, index, existing_caption=existing_caption) save_pic(splitted, index, existing_caption=existing_caption)
process_default_resize = False 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( autocrop_settings = autocrop.Settings(
crop_width = width, crop_width = width,
crop_height = height, crop_height = height,
face_points_weight = process_focal_crop_face_weight, face_points_weight = process_focal_crop_face_weight,
entropy_points_weight = process_focal_crop_entropy_weight, entropy_points_weight = process_focal_crop_entropy_weight,
corner_points_weight = process_focal_crop_edges_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): for focal in autocrop.crop_image(img, autocrop_settings):
save_pic(focal, index, existing_caption=existing_caption) save_pic(focal, index, existing_caption=existing_caption)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment