punimtag/scripts/analyze_poses.py

#!/usr/bin/env python3
"""
Analyze pose_mode values in the faces table
"""

import sqlite3
import sys
import os
from collections import Counter
from typing import Dict, List, Tuple

# Default database path
DEFAULT_DB_PATH = "data/photos.db"


def analyze_poses(db_path: str) -> None:
    """Analyze pose_mode values in faces table"""
    
    if not os.path.exists(db_path):
        print(f"❌ Database not found: {db_path}")
        return
    
    print(f"📊 Analyzing poses in database: {db_path}\n")
    
    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        # Get total number of faces
        cursor.execute("SELECT COUNT(*) FROM faces")
        total_faces = cursor.fetchone()[0]
        print(f"Total faces in database: {total_faces}\n")
        
        if total_faces == 0:
            print("No faces found in database.")
            conn.close()
            return
        
        # Get pose_mode distribution
        cursor.execute("""
            SELECT pose_mode, COUNT(*) as count
            FROM faces
            GROUP BY pose_mode
            ORDER BY count DESC
        """)
        
        pose_modes = cursor.fetchall()
        
        print("=" * 60)
        print("POSE_MODE DISTRIBUTION")
        print("=" * 60)
        for row in pose_modes:
            pose_mode = row['pose_mode'] or 'NULL'
            count = row['count']
            percentage = (count / total_faces) * 100
            print(f"  {pose_mode:30s} : {count:6d} ({percentage:5.1f}%)")
        
        print("\n" + "=" * 60)
        print("ANGLE STATISTICS")
        print("=" * 60)
        
        # Yaw angle statistics
        cursor.execute("""
            SELECT 
                COUNT(*) as total,
                COUNT(yaw_angle) as with_yaw,
                MIN(yaw_angle) as min_yaw,
                MAX(yaw_angle) as max_yaw,
                AVG(yaw_angle) as avg_yaw
            FROM faces
            WHERE yaw_angle IS NOT NULL
        """)
        yaw_stats = cursor.fetchone()
        
        # Pitch angle statistics
        cursor.execute("""
            SELECT 
                COUNT(*) as total,
                COUNT(pitch_angle) as with_pitch,
                MIN(pitch_angle) as min_pitch,
                MAX(pitch_angle) as max_pitch,
                AVG(pitch_angle) as avg_pitch
            FROM faces
            WHERE pitch_angle IS NOT NULL
        """)
        pitch_stats = cursor.fetchone()
        
        # Roll angle statistics
        cursor.execute("""
            SELECT 
                COUNT(*) as total,
                COUNT(roll_angle) as with_roll,
                MIN(roll_angle) as min_roll,
                MAX(roll_angle) as max_roll,
                AVG(roll_angle) as avg_roll
            FROM faces
            WHERE roll_angle IS NOT NULL
        """)
        roll_stats = cursor.fetchone()
        
        print(f"\nYaw Angle:")
        print(f"  Faces with yaw data: {yaw_stats['with_yaw']}")
        if yaw_stats['with_yaw'] > 0:
            print(f"  Min: {yaw_stats['min_yaw']:.1f}°")
            print(f"  Max: {yaw_stats['max_yaw']:.1f}°")
            print(f"  Avg: {yaw_stats['avg_yaw']:.1f}°")
        
        print(f"\nPitch Angle:")
        print(f"  Faces with pitch data: {pitch_stats['with_pitch']}")
        if pitch_stats['with_pitch'] > 0:
            print(f"  Min: {pitch_stats['min_pitch']:.1f}°")
            print(f"  Max: {pitch_stats['max_pitch']:.1f}°")
            print(f"  Avg: {pitch_stats['avg_pitch']:.1f}°")
        
        print(f"\nRoll Angle:")
        print(f"  Faces with roll data: {roll_stats['with_roll']}")
        if roll_stats['with_roll'] > 0:
            print(f"  Min: {roll_stats['min_roll']:.1f}°")
            print(f"  Max: {roll_stats['max_roll']:.1f}°")
            print(f"  Avg: {roll_stats['avg_roll']:.1f}°")
        
        # Sample faces with different poses
        print("\n" + "=" * 60)
        print("SAMPLE FACES BY POSE")
        print("=" * 60)
        
        for row in pose_modes[:10]:  # Top 10 pose modes
            pose_mode = row['pose_mode']
            cursor.execute("""
                SELECT id, photo_id, pose_mode, yaw_angle, pitch_angle, roll_angle
                FROM faces
                WHERE pose_mode = ?
                LIMIT 3
            """, (pose_mode,))
            samples = cursor.fetchall()
            
            print(f"\n{pose_mode}:")
            for sample in samples:
                yaw_str = f"{sample['yaw_angle']:.1f}°" if sample['yaw_angle'] is not None else "N/A"
                pitch_str = f"{sample['pitch_angle']:.1f}°" if sample['pitch_angle'] is not None else "N/A"
                roll_str = f"{sample['roll_angle']:.1f}°" if sample['roll_angle'] is not None else "N/A"
                print(f"  Face ID {sample['id']}: "
                      f"yaw={yaw_str} "
                      f"pitch={pitch_str} "
                      f"roll={roll_str}")
        
        conn.close()
        
    except sqlite3.Error as e:
        print(f"❌ Database error: {e}")
    except Exception as e:
        print(f"❌ Error: {e}")


def check_web_database() -> None:
    """Check if web database exists and analyze it"""
    # Common web database locations
    web_db_paths = [
        "data/punimtag.db",  # Default web database
        "data/web_photos.db",
        "data/photos_web.db",
        "web_photos.db",
    ]
    
    for db_path in web_db_paths:
        if os.path.exists(db_path):
            print(f"\n{'='*60}")
            print(f"WEB DATABASE: {db_path}")
            print(f"{'='*60}\n")
            analyze_poses(db_path)
            break


if __name__ == "__main__":
    # Check desktop database
    desktop_db = DEFAULT_DB_PATH
    if os.path.exists(desktop_db):
        analyze_poses(desktop_db)
    
    # Check web database
    check_web_database()
    
    # If no database found, list what we tried
    if not os.path.exists(desktop_db):
        print(f"❌ Desktop database not found: {desktop_db}")
        print("\nTrying to find database files...")
        for root, dirs, files in os.walk("data"):
            for file in files:
                if file.endswith(('.db', '.sqlite', '.sqlite3')):
                    print(f"  Found: {os.path.join(root, file)}")