punimtag/scripts/interactive_identifier.py

#!/usr/bin/env python3
"""
Interactive Face Identifier for PunimTag
Allows users to identify unknown faces in the database
"""

import os
import cv2
import numpy as np
from punimtag import PunimTag
from typing import Optional
import sys


class InteractiveFaceIdentifier:
    def __init__(self, db_path: str = 'punimtag.db'):
        self.tagger = PunimTag(db_path=db_path)
        self.window_name = 'Face Identifier'
        
    def display_face(self, image_path: str, location: tuple) -> np.ndarray:
        """Load and display image with face highlighted"""
        img = cv2.imread(image_path)
        if img is None:
            print(f"Error: Could not load image {image_path}")
            return None
            
        # Get face coordinates
        top, right, bottom, left = location
        
        # Draw rectangle around face
        cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 3)
        
        # Calculate display size (max 800x600)
        height, width = img.shape[:2]
        max_height, max_width = 600, 800
        
        if height > max_height or width > max_width:
            scale = min(max_height/height, max_width/width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height))
            
        return img
    
    def get_user_input(self, face_info: dict) -> Optional[str]:
        """Get user input for face identification"""
        print("\n" + "="*50)
        print(f"Image: {face_info['image_path']}")
        print(f"Face ID: {face_info['face_id']}")
        print("\nOptions:")
        print("1. Enter person's name")
        print("2. Skip this face (press Enter)")
        print("3. Quit (press 'q')")
        print("="*50)
        
        user_input = input("\nEnter person's name (or press Enter to skip): ").strip()
        
        if user_input.lower() == 'q':
            return 'QUIT'
        elif user_input == '':
            return None
        else:
            return user_input
    
    def run(self):
        """Run the interactive identification process"""
        print("PunimTag Interactive Face Identifier")
        print("=" * 50)
        
        # Get unidentified faces
        unidentified = self.tagger.get_unidentified_faces()
        
        if not unidentified:
            print("No unidentified faces found!")
            return
        
        print(f"Found {len(unidentified)} unidentified faces")
        
        # Create window
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        
        identified_count = 0
        
        for i, face_info in enumerate(unidentified):
            print(f"\nProcessing face {i+1} of {len(unidentified)}")
            
            # Display the face
            img = self.display_face(face_info['image_path'], face_info['location'])
            
            if img is None:
                continue
                
            cv2.imshow(self.window_name, img)
            cv2.waitKey(1)  # Allow window to update
            
            # Get user input
            name = self.get_user_input(face_info)
            
            if name == 'QUIT':
                print("\nQuitting...")
                break
            elif name:
                # Add person and assign face
                person_id = self.tagger.add_person(name)
                self.tagger.assign_face_to_person(face_info['face_id'], person_id, is_confirmed=True)
                identified_count += 1
                print(f"✓ Identified as: {name}")
            else:
                print("⊘ Skipped")
                
        cv2.destroyAllWindows()
        
        print(f"\n{'='*50}")
        print(f"Identification complete!")
        print(f"Identified {identified_count} faces")
        print(f"Skipped {len(unidentified) - identified_count} faces")
        
        self.tagger.close()


class CLIFaceIdentifier:
    """Command-line only face identifier (no OpenCV required)"""
    
    def __init__(self, db_path: str = 'punimtag.db'):
        self.tagger = PunimTag(db_path=db_path)
        
    def run(self):
        """Run CLI-based identification"""
        print("PunimTag CLI Face Identifier")
        print("=" * 50)
        
        # Get unidentified faces
        unidentified = self.tagger.get_unidentified_faces()
        
        if not unidentified:
            print("No unidentified faces found!")
            return
        
        print(f"Found {len(unidentified)} unidentified faces\n")
        
        identified_count = 0
        
        for i, face_info in enumerate(unidentified):
            print(f"\n{'='*50}")
            print(f"Face {i+1} of {len(unidentified)}")
            print(f"Image: {face_info['image_path']}")
            print(f"Location in image: top={face_info['location'][0]}, right={face_info['location'][1]}, "
                  f"bottom={face_info['location'][2]}, left={face_info['location'][3]}")
            print(f"Face ID: {face_info['face_id']}")
            
            name = input("\nEnter person's name (or press Enter to skip, 'q' to quit): ").strip()
            
            if name.lower() == 'q':
                print("\nQuitting...")
                break
            elif name:
                # Add person and assign face
                person_id = self.tagger.add_person(name)
                self.tagger.assign_face_to_person(face_info['face_id'], person_id, is_confirmed=True)
                identified_count += 1
                print(f"✓ Identified as: {name}")
            else:
                print("⊘ Skipped")
                
        print(f"\n{'='*50}")
        print(f"Identification complete!")
        print(f"Identified {identified_count} faces")
        print(f"Skipped {len(unidentified) - identified_count} faces")
        
        self.tagger.close()


def main():
    """Main entry point"""
    # Check if OpenCV is available
    try:
        import cv2
        print("OpenCV available - using visual identifier")
        identifier = InteractiveFaceIdentifier()
    except ImportError:
        print("OpenCV not available - using CLI identifier")
        print("Install opencv-python for visual identification: pip install opencv-python")
        identifier = CLIFaceIdentifier()
    
    try:
        identifier.run()
    except KeyboardInterrupt:
        print("\n\nInterrupted by user")
        sys.exit(0)


if __name__ == "__main__":
    main()