122 lines
2.6 KiB
Python
122 lines
2.6 KiB
Python
import sys
|
|
import json
|
|
import os
|
|
import piexif
|
|
import sqlite3
|
|
from sqlite3 import Error
|
|
from PIL import Image
|
|
import numpy as np
|
|
from deepface import DeepFace
|
|
from retinaface import RetinaFace
|
|
|
|
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()
|
|
|
|
model = DeepFace.build_model('ArcFace')
|
|
input_shape = DeepFace.functions.find_input_shape(model)
|
|
|
|
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')
|
|
faces = {}
|
|
identities = {}
|
|
|
|
with conn:
|
|
cur = conn.cursor()
|
|
res = cur.execute('''
|
|
SELECT faces.id,facedescriptors.descriptors
|
|
FROM faces
|
|
LEFT JOIN facedescriptors ON (faces.descriptorId=facedescriptors.id)
|
|
WHERE faces.identityId IS null
|
|
''')
|
|
for row in res.fetchall():
|
|
id, descriptors = row
|
|
if faces[id] is None:
|
|
face = {}
|
|
faces[id] = face
|
|
else:
|
|
face = faces[id]
|
|
face['descriptors'] = descriptors
|
|
|
|
# update_face_count(conn, photoId, len(faces))
|