punimtag/scripts/simple_identifier.py
2025-08-15 00:57:39 -08:00

118 lines
4.2 KiB
Python

#!/usr/bin/env python3
"""
Simple Face Identifier for PunimTag
Works with the punimtag_simple.db database
"""
import sqlite3
import os
from PIL import Image
import pickle
class SimpleFaceIdentifier:
def __init__(self, db_path='punimtag_simple.db'):
self.db_path = db_path
def get_unidentified_faces(self, limit=10):
"""Get a limited number of unidentified faces"""
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute('''SELECT f.id, f.image_id, i.path, i.filename, f.top, f.right, f.bottom, f.left
FROM faces f
JOIN images i ON f.image_id = i.id
WHERE f.person_id IS NULL
LIMIT ?''', (limit,))
faces = c.fetchall()
conn.close()
return faces
def add_person(self, name):
"""Add a new person"""
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute('INSERT OR IGNORE INTO people (name) VALUES (?)', (name,))
c.execute('SELECT id FROM people WHERE name = ?', (name,))
person_id = c.fetchone()[0]
conn.commit()
conn.close()
return person_id
def assign_face(self, face_id, person_id):
"""Assign a face to a person"""
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute('UPDATE faces SET person_id = ?, is_confirmed = 1 WHERE id = ?',
(person_id, face_id))
conn.commit()
conn.close()
def run_cli_identifier(self):
"""Run command line identifier"""
print("\n🏷️ Simple Face Identifier")
print("=" * 50)
faces = self.get_unidentified_faces(50) # Get first 50 faces
if not faces:
print("No unidentified faces found!")
return
print(f"Found {len(faces)} unidentified faces to process...")
print("For each face, enter the person's name or 's' to skip\n")
for i, (face_id, image_id, path, filename, top, right, bottom, left) in enumerate(faces):
print(f"\nFace {i+1}/{len(faces)}")
print(f"📁 File: {filename}")
print(f"📍 Location: top={top}, right={right}, bottom={bottom}, left={left}")
# Try to display basic info about the image
try:
if os.path.exists(path):
with Image.open(path) as img:
print(f"🖼️ Image size: {img.size}")
else:
print("⚠️ Image file not found")
except Exception as e:
print(f"⚠️ Could not read image: {e}")
while True:
name = input(f"👤 Who is this person? (or 's' to skip): ").strip()
if name.lower() == 's':
print("⏭️ Skipped")
break
elif name:
try:
person_id = self.add_person(name)
self.assign_face(face_id, person_id)
print(f"✅ Identified as '{name}'")
break
except Exception as e:
print(f"❌ Error: {e}")
else:
print("Please enter a name or 's' to skip")
print(f"\n🎉 Completed processing {len(faces)} faces!")
# Show remaining count
remaining = self.get_remaining_count()
if remaining > 0:
print(f"📊 {remaining} unidentified faces remaining")
print("Run the script again to continue identifying faces")
else:
print("🏆 All faces have been identified!")
def get_remaining_count(self):
"""Get count of remaining unidentified faces"""
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute('SELECT COUNT(*) FROM faces WHERE person_id IS NULL')
count = c.fetchone()[0]
conn.close()
return count
if __name__ == "__main__":
identifier = SimpleFaceIdentifier()
identifier.run_cli_identifier()