387 lines
8.7 KiB
Python
Executable file
387 lines
8.7 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import sqlite3
|
|
import sys
|
|
from collections import Counter, defaultdict
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from sklearn.cluster import HDBSCAN
|
|
|
|
|
|
EMBEDDING_DIMENSIONS = 128
|
|
EMBEDDING_BYTES = EMBEDDING_DIMENSIONS * 4
|
|
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Esegue HDBSCAN sugli embedding facciali "
|
|
"senza modificare il database."
|
|
),
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--database",
|
|
required=True,
|
|
help="Percorso del database SQLite",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--user",
|
|
required=True,
|
|
help="Utente proprietario dei volti",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--min-cluster-size",
|
|
type=int,
|
|
default=2,
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--min-samples",
|
|
type=int,
|
|
default=1,
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--metric",
|
|
default="euclidean",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--selection-method",
|
|
choices=("eom", "leaf"),
|
|
default="eom",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--scope",
|
|
choices=("all", "unclustered"),
|
|
default="all",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--pretty",
|
|
action="store_true",
|
|
)
|
|
|
|
arguments = parser.parse_args()
|
|
|
|
if arguments.min_cluster_size < 2:
|
|
parser.error(
|
|
"--min-cluster-size deve essere almeno 2"
|
|
)
|
|
|
|
if arguments.min_samples < 1:
|
|
parser.error(
|
|
"--min-samples deve essere almeno 1"
|
|
)
|
|
|
|
return arguments
|
|
|
|
|
|
def load_faces(database_path, scope, user):
|
|
connection = sqlite3.connect(database_path)
|
|
connection.row_factory = sqlite3.Row
|
|
|
|
query = """
|
|
SELECT
|
|
f.id,
|
|
f.photo_id,
|
|
f.face_index,
|
|
f.score,
|
|
f.embedding,
|
|
p.name AS photo_name,
|
|
p.path AS photo_path,
|
|
COALESCE(p.rotation, 0) AS photo_rotation
|
|
FROM faces AS f
|
|
JOIN photos AS p
|
|
ON p.id = f.photo_id
|
|
WHERE f.embedding IS NOT NULL
|
|
AND p.user = ?
|
|
"""
|
|
|
|
parameters = [user]
|
|
|
|
if scope == "unclustered":
|
|
query += """
|
|
AND f.cluster_id IS NULL
|
|
"""
|
|
|
|
query += """
|
|
ORDER BY
|
|
f.id
|
|
"""
|
|
|
|
rows = connection.execute(
|
|
query,
|
|
parameters,
|
|
).fetchall()
|
|
|
|
connection.close()
|
|
|
|
if len(rows) < 2:
|
|
raise RuntimeError(
|
|
"Sono necessari almeno due embedding"
|
|
)
|
|
|
|
faces = []
|
|
embeddings = []
|
|
|
|
for row in rows:
|
|
blob = row["embedding"]
|
|
|
|
if len(blob) != EMBEDDING_BYTES:
|
|
raise RuntimeError(
|
|
"Embedding non valido per "
|
|
f"{row['id']}: {len(blob)} byte"
|
|
)
|
|
|
|
vector = np.frombuffer(
|
|
blob,
|
|
dtype="<f4",
|
|
count=EMBEDDING_DIMENSIONS,
|
|
).astype(np.float64)
|
|
|
|
if not np.all(np.isfinite(vector)):
|
|
raise RuntimeError(
|
|
f"Embedding non finito: {row['id']}"
|
|
)
|
|
|
|
norm = float(np.linalg.norm(vector))
|
|
|
|
if not math.isfinite(norm) or norm <= 0.0:
|
|
raise RuntimeError(
|
|
f"Norma non valida: {row['id']}"
|
|
)
|
|
|
|
vector /= norm
|
|
|
|
faces.append(
|
|
{
|
|
"id": row["id"],
|
|
"photo_id": row["photo_id"],
|
|
"face_index": row["face_index"],
|
|
"detector_score": row["score"],
|
|
"photo_name": row["photo_name"],
|
|
"photo_path": row["photo_path"],
|
|
"photo_rotation": row["photo_rotation"],
|
|
"input_norm": norm,
|
|
}
|
|
)
|
|
|
|
embeddings.append(vector)
|
|
|
|
matrix = np.vstack(embeddings)
|
|
|
|
return faces, matrix
|
|
|
|
|
|
def cosine_statistics(indices, matrix):
|
|
if len(indices) < 2:
|
|
return {
|
|
"pair_count": 0,
|
|
"minimum": None,
|
|
"maximum": None,
|
|
"mean": None,
|
|
}
|
|
|
|
values = []
|
|
|
|
for left_position in range(len(indices)):
|
|
for right_position in range(
|
|
left_position + 1,
|
|
len(indices),
|
|
):
|
|
left_index = indices[left_position]
|
|
right_index = indices[right_position]
|
|
|
|
similarity = float(
|
|
np.dot(
|
|
matrix[left_index],
|
|
matrix[right_index],
|
|
)
|
|
)
|
|
|
|
values.append(similarity)
|
|
|
|
return {
|
|
"pair_count": len(values),
|
|
"minimum": float(min(values)),
|
|
"maximum": float(max(values)),
|
|
"mean": float(sum(values) / len(values)),
|
|
}
|
|
|
|
|
|
def build_result(arguments, faces, matrix, labels, probabilities):
|
|
label_counts = Counter(
|
|
int(label) for label in labels
|
|
)
|
|
|
|
cluster_indices = defaultdict(list)
|
|
|
|
for index, label_value in enumerate(labels):
|
|
label = int(label_value)
|
|
|
|
if label >= 0:
|
|
cluster_indices[label].append(index)
|
|
|
|
clusters = []
|
|
|
|
for label in sorted(cluster_indices):
|
|
indices = cluster_indices[label]
|
|
|
|
members = []
|
|
|
|
for index in indices:
|
|
face = dict(faces[index])
|
|
|
|
face["probability"] = float(
|
|
probabilities[index]
|
|
)
|
|
|
|
members.append(face)
|
|
|
|
clusters.append(
|
|
{
|
|
"label": label,
|
|
"size": len(indices),
|
|
"probability": {
|
|
"minimum": float(
|
|
min(
|
|
probabilities[index]
|
|
for index in indices
|
|
)
|
|
),
|
|
"maximum": float(
|
|
max(
|
|
probabilities[index]
|
|
for index in indices
|
|
)
|
|
),
|
|
"mean": float(
|
|
sum(
|
|
probabilities[index]
|
|
for index in indices
|
|
)
|
|
/ len(indices)
|
|
),
|
|
},
|
|
"cosine_similarity": cosine_statistics(
|
|
indices,
|
|
matrix,
|
|
),
|
|
"members": members,
|
|
}
|
|
)
|
|
|
|
noise_members = []
|
|
|
|
for index, label_value in enumerate(labels):
|
|
if int(label_value) != -1:
|
|
continue
|
|
|
|
face = dict(faces[index])
|
|
face["probability"] = float(
|
|
probabilities[index]
|
|
)
|
|
|
|
noise_members.append(face)
|
|
|
|
total_faces = len(faces)
|
|
noise_count = label_counts.get(-1, 0)
|
|
|
|
return {
|
|
"dry_run": True,
|
|
"database": str(arguments.database),
|
|
"parameters": {
|
|
"min_cluster_size":
|
|
arguments.min_cluster_size,
|
|
"min_samples":
|
|
arguments.min_samples,
|
|
"metric": arguments.metric,
|
|
"cluster_selection_method":
|
|
arguments.selection_method,
|
|
"scope": arguments.scope,
|
|
"user": arguments.user,
|
|
},
|
|
"summary": {
|
|
"total_faces": total_faces,
|
|
"cluster_count": len(clusters),
|
|
"clustered_faces":
|
|
total_faces - noise_count,
|
|
"noise_faces": noise_count,
|
|
"noise_percent": (
|
|
100.0 * noise_count / total_faces
|
|
),
|
|
},
|
|
"clusters": clusters,
|
|
"noise": noise_members,
|
|
}
|
|
|
|
|
|
def main():
|
|
arguments = parse_arguments()
|
|
|
|
database_path = Path(
|
|
arguments.database
|
|
).expanduser().resolve()
|
|
|
|
if not database_path.is_file():
|
|
raise RuntimeError(
|
|
f"Database non trovato: {database_path}"
|
|
)
|
|
|
|
faces, matrix = load_faces(
|
|
database_path,
|
|
arguments.scope,
|
|
arguments.user,
|
|
)
|
|
|
|
clusterer = HDBSCAN(
|
|
min_cluster_size=
|
|
arguments.min_cluster_size,
|
|
min_samples=arguments.min_samples,
|
|
metric=arguments.metric,
|
|
cluster_selection_method=
|
|
arguments.selection_method,
|
|
allow_single_cluster=False,
|
|
copy=True,
|
|
)
|
|
|
|
labels = clusterer.fit_predict(matrix)
|
|
probabilities = clusterer.probabilities_
|
|
|
|
result = build_result(
|
|
arguments,
|
|
faces,
|
|
matrix,
|
|
labels,
|
|
probabilities,
|
|
)
|
|
|
|
json.dump(
|
|
result,
|
|
sys.stdout,
|
|
indent=2 if arguments.pretty else None,
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
sys.stdout.write("\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception as error:
|
|
print(
|
|
f"[face_hdbscan] errore: {error}",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
raise
|