improve debug markers, fix algo weighting

1 files changed, 129 insertions, 78 deletions

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

+import os

 from collections import defaultdict

 from math import log, sqrt

 import numpy as np

@@ -26,19 +27,9 @@ def crop_image(im, settings):
       scale_by = settings.crop_height / im.height

 

   im = im.resize((int(im.width * scale_by), int(im.height * scale_by)))

+  im_debug = im.copy()

 

-  if im.width == settings.crop_width and im.height == settings.crop_height:

-    if settings.annotate_image:

-      d = ImageDraw.Draw(im)

-      rect = [0, 0, im.width, im.height]

-      rect[2] -= 1

-      rect[3] -= 1

-      d.rectangle(rect, outline=GREEN)

-      if settings.destop_view_image:

-        im.show()

-    return im

-

-  focus = focal_point(im, settings)

+  focus = focal_point(im_debug, settings)

 

   # take the focal point and turn it into crop coordinates that try to center over the focal

   # point but then get adjusted back into the frame

@@ -62,89 +53,143 @@ def crop_image(im, settings):
 

   crop = [x1, y1, x2, y2]

 

+  results = []

+

+  results.append(im.crop(tuple(crop)))

+

   if settings.annotate_image:

-    d = ImageDraw.Draw(im)

+    d = ImageDraw.Draw(im_debug)

     rect = list(crop)

     rect[2] -= 1

     rect[3] -= 1

     d.rectangle(rect, outline=GREEN)

+    results.append(im_debug)

     if settings.destop_view_image:

-      im.show()

+      im_debug.show()

 

-  return im.crop(tuple(crop))

+  return results

 

 def focal_point(im, settings):

     corner_points = image_corner_points(im, settings)

     entropy_points = image_entropy_points(im, settings)

     face_points = image_face_points(im, settings)

 

-    total_points = len(corner_points) + len(entropy_points) + len(face_points)

-

-    corner_weight = settings.corner_points_weight

-    entropy_weight = settings.entropy_points_weight

-    face_weight = settings.face_points_weight

-

-    weight_pref_total = corner_weight + entropy_weight + face_weight

-

-    # weight things

     pois = []

-    if weight_pref_total == 0 or total_points == 0: 

-      return pois

 

-    pois.extend(

-      [ PointOfInterest( p.x, p.y, weight=p.weight * ( (corner_weight/weight_pref_total) / (len(corner_points)/total_points) )) for p in corner_points ]

-    )

-    pois.extend(

-      [ PointOfInterest( p.x, p.y, weight=p.weight * ( (entropy_weight/weight_pref_total) / (len(entropy_points)/total_points) )) for p in entropy_points ]

-    )

-    pois.extend(

-      [ PointOfInterest( p.x, p.y, weight=p.weight * ( (face_weight/weight_pref_total) / (len(face_points)/total_points) )) for p in face_points ]

-    )

+    weight_pref_total = 0

+    if len(corner_points) > 0:

+      weight_pref_total += settings.corner_points_weight

+    if len(entropy_points) > 0:

+      weight_pref_total += settings.entropy_points_weight

+    if len(face_points) > 0:

+      weight_pref_total += settings.face_points_weight

+

+    corner_centroid = None

+    if len(corner_points) > 0:

+      corner_centroid = centroid(corner_points)

+      corner_centroid.weight = settings.corner_points_weight / weight_pref_total 

+      pois.append(corner_centroid)

+

+    entropy_centroid = None

+    if len(entropy_points) > 0:

+      entropy_centroid = centroid(entropy_points)

+      entropy_centroid.weight = settings.entropy_points_weight / weight_pref_total

+      pois.append(entropy_centroid)

+

+    face_centroid = None

+    if len(face_points) > 0:

+      face_centroid = centroid(face_points)

+      face_centroid.weight = settings.face_points_weight / weight_pref_total 

+      pois.append(face_centroid)

 

     average_point = poi_average(pois, settings)

 

     if settings.annotate_image:

       d = ImageDraw.Draw(im)

-      for f in face_points:

-        d.rectangle(f.bounding(f.size), outline=RED)

-      for f in entropy_points:

-        d.rectangle(f.bounding(30), outline=BLUE)

-      for poi in pois:

-        w = max(4, 4 * 0.5 * sqrt(poi.weight))

-        d.ellipse(poi.bounding(w), fill=BLUE)

-      d.ellipse(average_point.bounding(25), outline=GREEN)

+      max_size = min(im.width, im.height) * 0.07

+      if corner_centroid is not None:

+        color = BLUE

+        box = corner_centroid.bounding(max_size * corner_centroid.weight)

+        d.text((box[0], box[1]-15), "Edge: %.02f" % corner_centroid.weight, fill=color)

+        d.ellipse(box, outline=color)

+        if len(corner_points) > 1:

+          for f in corner_points:

+            d.rectangle(f.bounding(4), outline=color)

+      if entropy_centroid is not None:

+        color = "#ff0"

+        box = entropy_centroid.bounding(max_size * entropy_centroid.weight)

+        d.text((box[0], box[1]-15), "Entropy: %.02f" % entropy_centroid.weight, fill=color)

+        d.ellipse(box, outline=color)

+        if len(entropy_points) > 1:

+          for f in entropy_points:

+            d.rectangle(f.bounding(4), outline=color)

+      if face_centroid is not None:

+        color = RED

+        box = face_centroid.bounding(max_size * face_centroid.weight)

+        d.text((box[0], box[1]-15), "Face: %.02f" % face_centroid.weight, fill=color)

+        d.ellipse(box, outline=color)

+        if len(face_points) > 1:

+          for f in face_points:

+            d.rectangle(f.bounding(4), outline=color)

+

+      d.ellipse(average_point.bounding(max_size), outline=GREEN)

       

     return average_point

 

 

 def image_face_points(im, settings):

-    np_im = np.array(im)

-    gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY)

-

-    tries = [

-      [ f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01 ],

-      [ f'{cv2.data.haarcascades}haarcascade_frontalface_default.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_alt2.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_upperbody.xml', 0.05 ]

-    ]

-

-    for t in tries:

-      # print(t[0])

-      classifier = cv2.CascadeClassifier(t[0])

-      minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side

-      try:

-        faces = classifier.detectMultiScale(gray, scaleFactor=1.1,

-          minNeighbors=7, minSize=(minsize, minsize), flags=cv2.CASCADE_SCALE_IMAGE)

-      except:

-        continue

-

-      if len(faces) > 0:

-        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])) for r in rects]

+    if settings.dnn_model_path is not None:

+      detector = cv2.FaceDetectorYN.create(

+          settings.dnn_model_path,

+          "",

+          (im.width, im.height),

+          0.8, # score threshold

+          0.3, # nms threshold

+          5000 # keep top k before nms

+      )

+      faces = detector.detect(np.array(im))

+      results = []

+      if faces[1] is not None:

+        for face in faces[1]:

+          x = face[0]

+          y = face[1]

+          w = face[2]

+          h = face[3]

+          results.append(

+            PointOfInterest(

+              int(x + (w * 0.5)), # face focus left/right is center

+              int(y + (h * 0)), # face focus up/down is close to the top of the head

+              size = w,

+              weight = 1/len(faces[1])

+            )

+          )

+      return results

+    else:

+      np_im = np.array(im)

+      gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY)

+

+      tries = [

+        [ f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01 ],

+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_default.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_alt2.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_upperbody.xml', 0.05 ]

+      ]

+      for t in tries:

+        classifier = cv2.CascadeClassifier(t[0])

+        minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side

+        try:

+          faces = classifier.detectMultiScale(gray, scaleFactor=1.1,

+            minNeighbors=7, minSize=(minsize, minsize), flags=cv2.CASCADE_SCALE_IMAGE)

+        except:

+          continue

+

+        if len(faces) > 0:

+          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 []

 

 

@@ -161,7 +206,7 @@ def image_corner_points(im, settings):
         np_im,

         maxCorners=100,

         qualityLevel=0.04,

-        minDistance=min(grayscale.width, grayscale.height)*0.07,

+        minDistance=min(grayscale.width, grayscale.height)*0.03,

         useHarrisDetector=False,

     )

 

@@ -171,7 +216,7 @@ def image_corner_points(im, settings):
     focal_points = []

     for point in points:

       x, y = point.ravel()

-      focal_points.append(PointOfInterest(x, y, size=4))

+      focal_points.append(PointOfInterest(x, y, size=4, weight=1/len(points)))

 

     return focal_points

 

@@ -205,17 +250,22 @@ def image_entropy_points(im, settings):
     x_mid = int(crop_best[0] + settings.crop_width/2)

     y_mid = int(crop_best[1] + settings.crop_height/2)

 

-    return [PointOfInterest(x_mid, y_mid, size=25)]

+    return [PointOfInterest(x_mid, y_mid, size=25, weight=1.0)]

 

 

 def image_entropy(im):

     # greyscale image entropy

-    # band = np.asarray(im.convert("L"))

-    band = np.asarray(im.convert("1"), dtype=np.uint8)

+    band = np.asarray(im.convert("L"))

+    # band = np.asarray(im.convert("1"), dtype=np.uint8)

     hist, _ = np.histogram(band, bins=range(0, 256))

     hist = hist[hist > 0]

     return -np.log2(hist / hist.sum()).sum()

 

+def centroid(pois):

+  x = [poi.x for poi in pois]

+  y = [poi.y for poi in pois]

+  return PointOfInterest(sum(x)/len(pois), sum(y)/len(pois))

+

 

 def poi_average(pois, settings):

     weight = 0.0

@@ -260,11 +310,12 @@ class PointOfInterest:
 

 

 class Settings:

-  def __init__(self, crop_width=512, crop_height=512, corner_points_weight=0.5, entropy_points_weight=0.5, face_points_weight=0.5, annotate_image=False):

+  def __init__(self, crop_width=512, crop_height=512, corner_points_weight=0.5, entropy_points_weight=0.5, face_points_weight=0.5, annotate_image=False, dnn_model_path=None):

     self.crop_width = crop_width

     self.crop_height = crop_height

     self.corner_points_weight = corner_points_weight

     self.entropy_points_weight = entropy_points_weight

-    self.face_points_weight = entropy_points_weight

+    self.face_points_weight = face_points_weight

     self.annotate_image = annotate_image

-    self.destop_view_image = False
\ No newline at end of file
+    self.destop_view_image = False

+    self.dnn_model_path = dnn_model_path
\ No newline at end of file