Merge pull request #14121 from AUTOMATIC1111/fix-Auto-focal-point-crop-for-opencv-4.8.x

Fix auto focal point crop for opencv >= 4.8
This commit is contained in:
AUTOMATIC1111 2023-12-02 09:46:00 +03:00 committed by GitHub
commit 9eadc4f146
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 121 additions and 116 deletions

View File

@ -3,6 +3,8 @@ import requests
import os import os
import numpy as np import numpy as np
from PIL import ImageDraw from PIL import ImageDraw
from modules import paths_internal
from pkg_resources import parse_version
GREEN = "#0F0" GREEN = "#0F0"
BLUE = "#00F" BLUE = "#00F"
@ -25,7 +27,6 @@ def crop_image(im, settings):
elif is_portrait(settings.crop_width, settings.crop_height): elif is_portrait(settings.crop_width, settings.crop_height):
scale_by = settings.crop_height / im.height scale_by = settings.crop_height / im.height
im = im.resize((int(im.width * scale_by), int(im.height * scale_by))) im = im.resize((int(im.width * scale_by), int(im.height * scale_by)))
im_debug = im.copy() im_debug = im.copy()
@ -69,6 +70,7 @@ def crop_image(im, settings):
return results return results
def focal_point(im, settings): def focal_point(im, settings):
corner_points = image_corner_points(im, settings) if settings.corner_points_weight > 0 else [] corner_points = image_corner_points(im, settings) if settings.corner_points_weight > 0 else []
entropy_points = image_entropy_points(im, settings) if settings.entropy_points_weight > 0 else [] entropy_points = image_entropy_points(im, settings) if settings.entropy_points_weight > 0 else []
@ -78,118 +80,120 @@ def focal_point(im, settings):
weight_pref_total = 0 weight_pref_total = 0
if corner_points: if corner_points:
weight_pref_total += settings.corner_points_weight weight_pref_total += settings.corner_points_weight
if entropy_points: if entropy_points:
weight_pref_total += settings.entropy_points_weight weight_pref_total += settings.entropy_points_weight
if face_points: if face_points:
weight_pref_total += settings.face_points_weight weight_pref_total += settings.face_points_weight
corner_centroid = None corner_centroid = None
if corner_points: if corner_points:
corner_centroid = centroid(corner_points) corner_centroid = centroid(corner_points)
corner_centroid.weight = settings.corner_points_weight / weight_pref_total corner_centroid.weight = settings.corner_points_weight / weight_pref_total
pois.append(corner_centroid) pois.append(corner_centroid)
entropy_centroid = None entropy_centroid = None
if entropy_points: if entropy_points:
entropy_centroid = centroid(entropy_points) entropy_centroid = centroid(entropy_points)
entropy_centroid.weight = settings.entropy_points_weight / weight_pref_total entropy_centroid.weight = settings.entropy_points_weight / weight_pref_total
pois.append(entropy_centroid) pois.append(entropy_centroid)
face_centroid = None face_centroid = None
if face_points: if face_points:
face_centroid = centroid(face_points) face_centroid = centroid(face_points)
face_centroid.weight = settings.face_points_weight / weight_pref_total face_centroid.weight = settings.face_points_weight / weight_pref_total
pois.append(face_centroid) pois.append(face_centroid)
average_point = poi_average(pois, settings) average_point = poi_average(pois, settings)
if settings.annotate_image: if settings.annotate_image:
d = ImageDraw.Draw(im) d = ImageDraw.Draw(im)
max_size = min(im.width, im.height) * 0.07 max_size = min(im.width, im.height) * 0.07
if corner_centroid is not None: if corner_centroid is not None:
color = BLUE color = BLUE
box = corner_centroid.bounding(max_size * corner_centroid.weight) box = corner_centroid.bounding(max_size * corner_centroid.weight)
d.text((box[0], box[1]-15), f"Edge: {corner_centroid.weight:.02f}", fill=color) d.text((box[0], box[1] - 15), f"Edge: {corner_centroid.weight:.02f}", fill=color)
d.ellipse(box, outline=color) d.ellipse(box, outline=color)
if len(corner_points) > 1: if len(corner_points) > 1:
for f in corner_points: for f in corner_points:
d.rectangle(f.bounding(4), outline=color) d.rectangle(f.bounding(4), outline=color)
if entropy_centroid is not None: if entropy_centroid is not None:
color = "#ff0" color = "#ff0"
box = entropy_centroid.bounding(max_size * entropy_centroid.weight) box = entropy_centroid.bounding(max_size * entropy_centroid.weight)
d.text((box[0], box[1]-15), f"Entropy: {entropy_centroid.weight:.02f}", fill=color) d.text((box[0], box[1] - 15), f"Entropy: {entropy_centroid.weight:.02f}", fill=color)
d.ellipse(box, outline=color) d.ellipse(box, outline=color)
if len(entropy_points) > 1: if len(entropy_points) > 1:
for f in entropy_points: for f in entropy_points:
d.rectangle(f.bounding(4), outline=color) d.rectangle(f.bounding(4), outline=color)
if face_centroid is not None: if face_centroid is not None:
color = RED color = RED
box = face_centroid.bounding(max_size * face_centroid.weight) box = face_centroid.bounding(max_size * face_centroid.weight)
d.text((box[0], box[1]-15), f"Face: {face_centroid.weight:.02f}", fill=color) d.text((box[0], box[1] - 15), f"Face: {face_centroid.weight:.02f}", fill=color)
d.ellipse(box, outline=color) d.ellipse(box, outline=color)
if len(face_points) > 1: if len(face_points) > 1:
for f in face_points: for f in face_points:
d.rectangle(f.bounding(4), outline=color) d.rectangle(f.bounding(4), outline=color)
d.ellipse(average_point.bounding(max_size), outline=GREEN) d.ellipse(average_point.bounding(max_size), outline=GREEN)
return average_point return average_point
def image_face_points(im, settings): def image_face_points(im, settings):
if settings.dnn_model_path is not None: if settings.dnn_model_path is not None:
detector = cv2.FaceDetectorYN.create( detector = cv2.FaceDetectorYN.create(
settings.dnn_model_path, settings.dnn_model_path,
"", "",
(im.width, im.height), (im.width, im.height),
0.9, # score threshold 0.9, # score threshold
0.3, # nms threshold 0.3, # nms threshold
5000 # keep top k before nms 5000 # keep top k before nms
) )
faces = detector.detect(np.array(im)) faces = detector.detect(np.array(im))
results = [] results = []
if faces[1] is not None: if faces[1] is not None:
for face in faces[1]: for face in faces[1]:
x = face[0] x = face[0]
y = face[1] y = face[1]
w = face[2] w = face[2]
h = face[3] h = face[3]
results.append( results.append(
PointOfInterest( PointOfInterest(
int(x + (w * 0.5)), # face focus left/right is center int(x + (w * 0.5)), # face focus left/right is center
int(y + (h * 0.33)), # face focus up/down is close to the top of the head int(y + (h * 0.33)), # face focus up/down is close to the top of the head
size = w, size=w,
weight = 1/len(faces[1]) weight=1 / len(faces[1])
) )
) )
return results return results
else: else:
np_im = np.array(im) np_im = np.array(im)
gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY)
tries = [ tries = [
[ f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01 ], [f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01],
[ f'{cv2.data.haarcascades}haarcascade_frontalface_default.xml', 0.05 ], [f'{cv2.data.haarcascades}haarcascade_frontalface_default.xml', 0.05],
[ f'{cv2.data.haarcascades}haarcascade_profileface.xml', 0.05 ], [f'{cv2.data.haarcascades}haarcascade_profileface.xml', 0.05],
[ f'{cv2.data.haarcascades}haarcascade_frontalface_alt.xml', 0.05 ], [f'{cv2.data.haarcascades}haarcascade_frontalface_alt.xml', 0.05],
[ f'{cv2.data.haarcascades}haarcascade_frontalface_alt2.xml', 0.05 ], [f'{cv2.data.haarcascades}haarcascade_frontalface_alt2.xml', 0.05],
[ f'{cv2.data.haarcascades}haarcascade_frontalface_alt_tree.xml', 0.05 ], [f'{cv2.data.haarcascades}haarcascade_frontalface_alt_tree.xml', 0.05],
[ f'{cv2.data.haarcascades}haarcascade_eye_tree_eyeglasses.xml', 0.05 ], [f'{cv2.data.haarcascades}haarcascade_eye_tree_eyeglasses.xml', 0.05],
[ f'{cv2.data.haarcascades}haarcascade_upperbody.xml', 0.05 ] [f'{cv2.data.haarcascades}haarcascade_upperbody.xml', 0.05]
] ]
for t in tries: for t in tries:
classifier = cv2.CascadeClassifier(t[0]) classifier = cv2.CascadeClassifier(t[0])
minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side
try: try:
faces = classifier.detectMultiScale(gray, scaleFactor=1.1, faces = classifier.detectMultiScale(gray, scaleFactor=1.1,
minNeighbors=7, minSize=(minsize, minsize), flags=cv2.CASCADE_SCALE_IMAGE) minNeighbors=7, minSize=(minsize, minsize),
except Exception: flags=cv2.CASCADE_SCALE_IMAGE)
continue except Exception:
continue
if faces: if faces:
rects = [[f[0], f[1], f[0] + f[2], f[1] + f[3]] for f in faces] rects = [[f[0], f[1], f[0] + f[2], f[1] + f[3]] for f in faces]
return [PointOfInterest((r[0] +r[2]) // 2, (r[1] + r[3]) // 2, size=abs(r[0]-r[2]), weight=1/len(rects)) for r in rects] return [PointOfInterest((r[0] + r[2]) // 2, (r[1] + r[3]) // 2, size=abs(r[0] - r[2]),
weight=1 / len(rects)) for r in rects]
return [] return []
@ -198,7 +202,7 @@ def image_corner_points(im, settings):
# naive attempt at preventing focal points from collecting at watermarks near the bottom # naive attempt at preventing focal points from collecting at watermarks near the bottom
gd = ImageDraw.Draw(grayscale) gd = ImageDraw.Draw(grayscale)
gd.rectangle([0, im.height*.9, im.width, im.height], fill="#999") gd.rectangle([0, im.height * .9, im.width, im.height], fill="#999")
np_im = np.array(grayscale) np_im = np.array(grayscale)
@ -206,7 +210,7 @@ def image_corner_points(im, settings):
np_im, np_im,
maxCorners=100, maxCorners=100,
qualityLevel=0.04, qualityLevel=0.04,
minDistance=min(grayscale.width, grayscale.height)*0.06, minDistance=min(grayscale.width, grayscale.height) * 0.06,
useHarrisDetector=False, useHarrisDetector=False,
) )
@ -215,8 +219,8 @@ def image_corner_points(im, settings):
focal_points = [] focal_points = []
for point in points: for point in points:
x, y = point.ravel() x, y = point.ravel()
focal_points.append(PointOfInterest(x, y, size=4, weight=1/len(points))) focal_points.append(PointOfInterest(x, y, size=4, weight=1 / len(points)))
return focal_points return focal_points
@ -225,13 +229,13 @@ def image_entropy_points(im, settings):
landscape = im.height < im.width landscape = im.height < im.width
portrait = im.height > im.width portrait = im.height > im.width
if landscape: if landscape:
move_idx = [0, 2] move_idx = [0, 2]
move_max = im.size[0] move_max = im.size[0]
elif portrait: elif portrait:
move_idx = [1, 3] move_idx = [1, 3]
move_max = im.size[1] move_max = im.size[1]
else: else:
return [] return []
e_max = 0 e_max = 0
crop_current = [0, 0, settings.crop_width, settings.crop_height] crop_current = [0, 0, settings.crop_width, settings.crop_height]
@ -241,14 +245,14 @@ def image_entropy_points(im, settings):
e = image_entropy(crop) e = image_entropy(crop)
if (e > e_max): if (e > e_max):
e_max = e e_max = e
crop_best = list(crop_current) crop_best = list(crop_current)
crop_current[move_idx[0]] += 4 crop_current[move_idx[0]] += 4
crop_current[move_idx[1]] += 4 crop_current[move_idx[1]] += 4
x_mid = int(crop_best[0] + settings.crop_width/2) x_mid = int(crop_best[0] + settings.crop_width / 2)
y_mid = int(crop_best[1] + settings.crop_height/2) y_mid = int(crop_best[1] + settings.crop_height / 2)
return [PointOfInterest(x_mid, y_mid, size=25, weight=1.0)] return [PointOfInterest(x_mid, y_mid, size=25, weight=1.0)]
@ -294,22 +298,23 @@ def is_square(w, h):
return w == h return w == h
def download_and_cache_models(dirname): model_dir_opencv = os.path.join(paths_internal.models_path, 'opencv')
download_url = 'https://github.com/opencv/opencv_zoo/blob/91fb0290f50896f38a0ab1e558b74b16bc009428/models/face_detection_yunet/face_detection_yunet_2022mar.onnx?raw=true' if parse_version(cv2.__version__) >= parse_version('4.8'):
model_file_name = 'face_detection_yunet.onnx' model_file_path = os.path.join(model_dir_opencv, 'face_detection_yunet_2023mar.onnx')
model_url = 'https://github.com/opencv/opencv_zoo/blob/b6e370b10f641879a87890d44e42173077154a05/models/face_detection_yunet/face_detection_yunet_2023mar.onnx?raw=true'
else:
model_file_path = os.path.join(model_dir_opencv, 'face_detection_yunet.onnx')
model_url = 'https://github.com/opencv/opencv_zoo/blob/91fb0290f50896f38a0ab1e558b74b16bc009428/models/face_detection_yunet/face_detection_yunet_2022mar.onnx?raw=true'
os.makedirs(dirname, exist_ok=True)
cache_file = os.path.join(dirname, model_file_name) def download_and_cache_models():
if not os.path.exists(cache_file): if not os.path.exists(model_file_path):
print(f"downloading face detection model from '{download_url}' to '{cache_file}'") os.makedirs(model_dir_opencv, exist_ok=True)
response = requests.get(download_url) print(f"downloading face detection model from '{model_url}' to '{model_file_path}'")
with open(cache_file, "wb") as f: response = requests.get(model_url)
with open(model_file_path, "wb") as f:
f.write(response.content) f.write(response.content)
return model_file_path
if os.path.exists(cache_file):
return cache_file
return None
class PointOfInterest: class PointOfInterest:

View File

@ -3,7 +3,7 @@ from PIL import Image, ImageOps
import math import math
import tqdm import tqdm
from modules import paths, shared, images, deepbooru from modules import shared, images, deepbooru
from modules.textual_inversion import autocrop from modules.textual_inversion import autocrop
@ -196,7 +196,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
dnn_model_path = None dnn_model_path = None
try: try:
dnn_model_path = autocrop.download_and_cache_models(os.path.join(paths.models_path, "opencv")) dnn_model_path = autocrop.download_and_cache_models()
except Exception as e: except Exception as e:
print("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", e) print("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", e)