ketr.photos/server/detect.py

import sys
import json
import os
import piexif
import sqlite3
from sqlite3 import Error
from PIL import Image
from deepface import DeepFace
from retinaface import RetinaFace
import numpy as np

class NpEncoder(json.JSONEncoder):
  def default(self, obj):
    if isinstance(obj, np.integer):
      return int(obj)
    if isinstance(obj, np.floating):
      return float(obj)
    if isinstance(obj, np.ndarray):
      return obj.tolist()
        
models = ["VGG-Face", "Facenet", "Facenet512", "OpenFace", "DeepFace", "DeepID", "ArcFace", "Dlib", "SFace"]

model = DeepFace.build_model('ArcFace')
input_shape = DeepFace.functions.find_input_shape(model)


# Adapted from DeepFace
# https://github.com/serengil/deepface/blob/master/deepface/commons/functions.py
#
# Modified to use bicubic resampling and clip expansion, as well as to
# take a PIL Image instead of numpy array
def alignment_procedure(img, left_eye, right_eye):
  """
  Given left and right eye coordinates in image, rotate around point
  between eyes such that eyes are horizontal
  :param img: Image (not np.array)
  :param left_eye: Eye appearing on the left (right eye of person)
  :param right_eye: Eye appearing on the right (left eye of person)
  :return: adjusted image
  """
  dY = right_eye[1] - left_eye[1]
  dX = right_eye[0] - left_eye[0]
  rotation = -np.atan2(dY, dX)
#  cosRotation = np.cos(rotation)
#  sinRotation = np.sin(rotation)
#  eyeDistance = np.sqrt(dY * dY + dX * dX)
#  mid_x = left_eye[0] + 0.5 * dX
#  mid_y = left_eye[1] + 0.5 * dY
#  prime_x = mid_x * cosRotation - mid_y * sinRotation
#  prime_y = mid_y * cosRotation - mid_x * sinRotation

  img = img.rotate(
    angle = np.pi * rotation,
    resample=Image.BICUBIC, 
    expand=True)
  
  return img

def extract_faces(img, threshold=0.9, model = None, allow_upscaling = True):
  faces = RetinaFace.detect_faces(img_path = img, threshold = threshold, model = model, allow_upscaling = allow_upscaling)
  #faces = DeepFace.detectFace(img_path = img, target_size = (224, 224),    detector_backend = 'retinaface')

  # Re-implementation of 'extract_faces' with the addition of keeping a 
  # copy of the face image for caching on disk
  if type(faces) == dict:
    for key in faces:
      identity = faces[key]
      facial_area = identity["facial_area"]

      width = facial_area[2] - facial_area[0]
      height = facial_area[3] - facial_area[1]
      x = facial_area[0] + width * 0.5
      y = facial_area[1] + height * 0.5

      # Make thumbnail a square crop
      if width > height:
        height = width
      else:
        width = height

      landmarks = identity["landmarks"]
      left_eye = landmarks["left_eye"]
      right_eye = landmarks["right_eye"]
      nose = landmarks["nose"]

      # translate the landmarks to be centered on array
      left_eye[0] -= x
      left_eye[1] -= y
      right_eye[0] -= x
      right_eye[1] -= y
      nose[0] -= x
      nose[1] -= y

      width *= 1.25
      height *= 1.25

      left = max(round(x - width * 0.5), facial_area[0])
      right = min(round(left + width), facial_area[2])
      top = max(round(y - height * 0.5), facial_area[1])
      bottom = min(round(top + height), facial_area[3])

      facial_img = img[top: bottom, left: right]

      # Eye order is reversed as the routine does them backwards
      aligned = RetinaFace.postprocess.alignment_procedure(facial_img, right_eye, left_eye, nose)

      image = Image.fromarray(aligned)
      image = image.resize(size = input_shape, resample = Image.LANCZOS)
      resized = np.asarray(image)
        
      identity['vector'] = DeepFace.represent(
        img_path = resized,
        model_name = 'ArcFace',
        model = model, # pre-built
        detector_backend = 'retinaface',
        enforce_detection = False)

      identity["face"] = { 
        'top': facial_area[1] / img.shape[0],
        'left': facial_area[0] / img.shape[1],
        'bottom': facial_area[3] / img.shape[0],
        'right': facial_area[2] / img.shape[1]
      }

      identity['image'] = resized #[:, :, ::-1]

    return faces

#face verification
#img_path = sys.argv[1]

def create_connection(db_file):
  """ create a database connection to the SQLite database
      specified by db_file
  :param db_file: database file
  :return: Connection object or None
  """
  conn = None
  try:
    conn = sqlite3.connect(db_file)
  except Error as e:
    print(e)

  return conn

def create_face(conn, face):
  """
  Create a new face in the faces table
  :param conn:
  :param face:
  :return: face id
  """
  sql = '''
  INSERT INTO faces(photoId,scanVersion,faceConfidence,top,left,bottom,right)
  VALUES(?,?,?,?,?,?,?)
  '''
  cur = conn.cursor()
  cur.execute(sql, (
    face['photoId'], 
    face['scanVersion'], 
    face['faceConfidence'],
    face['top'],
    face['left'],
    face['bottom'],
    face['right']
  ))
  conn.commit()
  return cur.lastrowid

def create_face_descriptor(conn, faceId, descriptor):
  """
  Create a new face in the faces table
  :param conn:
  :param faceId:
  :param descriptor:
  :return: descriptor id
  """
  sql = '''
  INSERT INTO facedescriptors(faceId,model,descriptors)
  VALUES(?,?,?)
  '''
  cur = conn.cursor()
  cur.execute(sql, (
    faceId, 
    descriptor['model'],
    np.array(descriptor['descriptors'])
  ))
  conn.commit()
  return cur.lastrowid

def update_face_count(conn, photoId, faces):
  """
  Update the number of faces that have been matched on a photo
  :param conn:
  :param photoId:
  :param faces:
  :return: None
  """
  sql = '''
  UPDATE photos SET faces=? WHERE id=?
  '''
  cur = conn.cursor()
  cur.execute(sql, (faces, photoId))
  conn.commit()
  return None


base = '/pictures/'
conn = create_connection('../db/photos.db')
with conn:
  cur = conn.cursor()
  for row in cur.execute('''
    SELECT photos.id,photos.faces,albums.path,photos.filename FROM photos
    LEFT JOIN albums ON (albums.id=photos.albumId) 
    WHERE photos.faces=-1
    '''):
    photoId, photoFaces, albumPath, photoFilename = row
    img_path = f'{base}{albumPath}{photoFilename}'
    print(f'Processing {img_path}')
    img = Image.open(img_path)
    img = img.convert()
    img = np.asarray(img)
    print(img.shape)
    faces = extract_faces(img)
    if faces is None:
      update_face_count(conn, photoId, 0)
      continue
    print(f'Handling {len(faces)} faces')
    for key in faces:
      face = faces[key]
      image = Image.fromarray(face['image'])
      
      #face['analysis'] = DeepFace.analyze(img_path = img, actions = ['age', 'gender', 'race', 'emotion'], enforce_detection = False)
      #face['analysis'] = DeepFace.analyze(img, actions = ['emotion'])

      # TODO: Add additional meta-data allowing back referencing to original
      # photo
      face['version'] = 1 # version 1 doesn't add much...

      data = {k: face[k] for k in set(list(face.keys())) - set(['image', 'facial_area', 'landmarks'])}
      json_str = json.dumps(data, ensure_ascii=False, indent=2, cls=NpEncoder)

      faceId = create_face(conn, {
        'photoId': photoId,
        'scanVersion': face['version'],
        'faceConfidence': face['score'],
        'top': face['face']['top'],
        'left': face['face']['left'],
        'bottom': face['face']['bottom'],
        'right': face['face']['right'],
      })

      faceDescriptorId = create_face_descriptor(conn, faceId, {
        'model': 'RetinaFace',
        'descriptors': face['vector']
      })

      path = f'faces/{faceId % 100}'
      try:
        os.mkdir(path)
      except FileExistsError:
        pass
        
      with open(f'{path}/{faceId}.json', 'w', encoding = 'utf-8') as f:
        f.write(json_str)

      # Encode this data into the JPG as Exif
      exif_ifd = {piexif.ExifIFD.UserComment: json_str.encode()}
      exif_dict = {"0th": {}, "Exif": exif_ifd, "1st": {},
        "thumbnail": None, "GPS": {}}
      image.save(f'{path}/{faceId}.jpg', exif = piexif.dump(exif_dict))

      #df = DeepFace.find(img, db_path = '/db')
      #print(df.head())

    update_face_count(conn, photoId, len(faces))

  #img2_path = sys.argv[2]
#print("image 1: ", img1_path);
#print("image 2: ", img2_path);
#result = DeepFace.verify(img1_path = img1_path, img2_path = img2_path, #model_name = models[1])
#print("result: ", result)

#face recognition
#df = DeepFace.find(img_path = img1_path, db_path = "./db/deepface", model_name = models[1])

#print("df: ", df)