Commit 3e6c2420 authored by captin411's avatar captin411
Browse files

improve debug markers, fix algo weighting

parent 1be5933b
Loading
Loading
Loading
Loading
+129 −78
Original line number Diff line number Diff line
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,62 +53,118 @@ 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)
      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(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)
            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):
    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)

@@ -131,9 +178,7 @@ def image_face_points(im, settings):
        [ 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:
@@ -144,7 +189,7 @@ def image_face_points(im, settings):

        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]
          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
    self.dnn_model_path = dnn_model_path
 No newline at end of file