Add punimtag.py script and update requirements.txt
This commit is contained in:
parent
9cf63b887b
commit
79284d0a5c
94
punimtag.py
Normal file
94
punimtag.py
Normal file
@ -0,0 +1,94 @@
|
||||
import os
|
||||
import sqlite3
|
||||
import face_recognition
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import hdbscan
|
||||
|
||||
# Constants
|
||||
PHOTOS_DIR = 'photos'
|
||||
DB_FILE = 'faces.db'
|
||||
|
||||
# Create or connect to database
|
||||
def init_db():
|
||||
conn = sqlite3.connect(DB_FILE)
|
||||
c = conn.cursor()
|
||||
c.execute('''CREATE TABLE IF NOT EXISTS images
|
||||
(id INTEGER PRIMARY KEY, path TEXT UNIQUE)''')
|
||||
c.execute('''CREATE TABLE IF NOT EXISTS faces
|
||||
(id INTEGER PRIMARY KEY, image_id INTEGER, location TEXT, encoding BLOB, cluster_id INTEGER,
|
||||
FOREIGN KEY(image_id) REFERENCES images(id))''')
|
||||
c.execute('''CREATE TABLE IF NOT EXISTS clusters
|
||||
(id INTEGER PRIMARY KEY, label TEXT)''')
|
||||
conn.commit()
|
||||
return conn
|
||||
|
||||
# Walk folder and process images
|
||||
def process_images(conn):
|
||||
c = conn.cursor()
|
||||
encodings = []
|
||||
face_data = []
|
||||
|
||||
for root, _, files in os.walk(PHOTOS_DIR):
|
||||
for file in files:
|
||||
if file.lower().endswith(('.jpg', '.png')):
|
||||
path = os.path.join(root, file)
|
||||
try:
|
||||
image = face_recognition.load_image_file(path)
|
||||
locations = face_recognition.face_locations(image)
|
||||
face_encodings = face_recognition.face_encodings(image, locations)
|
||||
|
||||
c.execute("INSERT OR IGNORE INTO images (path) VALUES (?)", (path,))
|
||||
c.execute("SELECT id FROM images WHERE path = ?", (path,))
|
||||
image_id = c.fetchone()[0]
|
||||
|
||||
for loc, enc in zip(locations, face_encodings):
|
||||
loc_str = str(loc)
|
||||
enc_blob = enc.tobytes()
|
||||
encodings.append(enc)
|
||||
face_data.append((image_id, loc_str, enc_blob))
|
||||
except Exception as e:
|
||||
print(f"Error processing {path}: {e}")
|
||||
|
||||
conn.commit()
|
||||
return np.array(encodings), face_data
|
||||
|
||||
# Cluster encodings
|
||||
def cluster_encodings(encodings):
|
||||
if len(encodings) == 0:
|
||||
return []
|
||||
clusterer = hdbscan.HDBSCAN(min_cluster_size=2)
|
||||
labels = clusterer.fit_predict(encodings)
|
||||
return labels
|
||||
|
||||
# Store clusters and update faces
|
||||
def store_clusters(conn, labels, face_data):
|
||||
c = conn.cursor()
|
||||
unique_labels = set(labels)
|
||||
for label in unique_labels:
|
||||
if label != -1: # Ignore noise
|
||||
c.execute("INSERT INTO clusters (label) VALUES (?)", (f"Cluster {label}",))
|
||||
c.execute("SELECT last_insert_rowid()")
|
||||
cluster_id = c.fetchone()[0]
|
||||
else:
|
||||
cluster_id = None
|
||||
|
||||
# Update faces with this label
|
||||
for i, (image_id, loc_str, enc_blob) in enumerate(face_data):
|
||||
if labels[i] == label:
|
||||
c.execute("INSERT INTO faces (image_id, location, encoding, cluster_id) VALUES (?, ?, ?, ?)",
|
||||
(image_id, loc_str, enc_blob, cluster_id))
|
||||
|
||||
conn.commit()
|
||||
|
||||
# Main function
|
||||
def main():
|
||||
conn = init_db()
|
||||
encodings, face_data = process_images(conn)
|
||||
labels = cluster_encodings(encodings)
|
||||
store_clusters(conn, labels, face_data)
|
||||
conn.close()
|
||||
print("Processing complete. Data stored in faces.db")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -2,5 +2,10 @@ click==8.2.1
|
||||
dlib==20.0.0
|
||||
face-recognition==1.3.0
|
||||
face-recognition-models==0.3.0
|
||||
hdbscan==0.8.40
|
||||
joblib==1.5.1
|
||||
numpy==2.2.6
|
||||
pillow==11.3.0
|
||||
scikit-learn==1.7.0
|
||||
scipy==1.15.3
|
||||
threadpoolctl==3.6.0
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user