image.py

import os
import face_recognition as fr
import pickle
import shutil

# Use persistent storage path
DATA_DIR = "/data"
KNOWN_USER_DIR = os.path.join(DATA_DIR, "known_users")
UNKNOWN_USER_DIR = os.path.join(DATA_DIR, "unknown_users")
ENCODINGS_DIR = os.path.join(DATA_DIR, "encodings")

# Ensure directories exist
os.makedirs(KNOWN_USER_DIR, exist_ok=True)
os.makedirs(UNKNOWN_USER_DIR, exist_ok=True)
os.makedirs(ENCODINGS_DIR, exist_ok=True)


def save_image_from_path(image_path, student_id):
    """Save student image from file path and compute face encoding"""
    try:
        if not image_path or not os.path.exists(image_path):
            return False, "No image provided"
        
        # Destination path
        dest_path = os.path.join(KNOWN_USER_DIR, f"{student_id}.jpg")
        
        # Copy image to destination
        shutil.copy2(image_path, dest_path)
        
        # Compute and save face encoding
        image = fr.load_image_file(dest_path)
        face_encodings = fr.face_encodings(image)
        
        if len(face_encodings) > 0:
            encoding = face_encodings[0]
            encoding_path = os.path.join(ENCODINGS_DIR, f"{student_id}.pkl")
            
            with open(encoding_path, "wb") as f:
                pickle.dump(encoding, f)
            
            return True, "Image and face encoding saved successfully!"
        else:
            # Delete image if no face found
            if os.path.exists(dest_path):
                os.remove(dest_path)
            return False, "No face detected in the image. Please upload a clear face photo."
            
    except Exception as e:
        return False, f"Error processing image: {str(e)}"


def save_temp_image_from_path(image_path):
    """Save temporary image from file path for login"""
    try:
        if not image_path or not os.path.exists(image_path):
            return None
        
        temp_path = os.path.join(UNKNOWN_USER_DIR, "temp_login.jpg")
        shutil.copy2(image_path, temp_path)
        return temp_path
        
    except Exception as e:
        print(f"Error saving temp image: {e}")
        return None


def delete_temp_image():
    """Delete temporary login image"""
    temp_path = os.path.join(UNKNOWN_USER_DIR, "temp_login.jpg")
    if os.path.exists(temp_path):
        try:
            os.remove(temp_path)
        except Exception as e:
            print(f"Error deleting temp image: {e}")


def recognize_face(unknown_image_path, tolerance=0.5):
    """
    Recognize face by comparing with stored encodings
    Returns: (is_match, student_id, confidence, message)
    """
    try:
        # Load unknown image
        if not os.path.exists(unknown_image_path):
            return False, None, 0, "Image file not found"
        
        unknown_image = fr.load_image_file(unknown_image_path)
        unknown_encodings = fr.face_encodings(unknown_image)
        
        if len(unknown_encodings) == 0:
            return False, None, 0, "No face detected in the image"
        
        unknown_encoding = unknown_encodings[0]
        
        # Compare with all stored encodings
        best_match_id = None
        best_distance = float('inf')
        
        encoding_files = [f for f in os.listdir(ENCODINGS_DIR) if f.endswith('.pkl')]
        
        if not encoding_files:
            return False, None, 0, "No registered students found in database"
        
        for encoding_file in encoding_files:
            student_id = encoding_file.replace('.pkl', '')
            encoding_path = os.path.join(ENCODINGS_DIR, encoding_file)
            
            with open(encoding_path, 'rb') as f:
                known_encoding = pickle.load(f)
            
            distance = fr.face_distance([known_encoding], unknown_encoding)[0]
            
            if distance < best_distance:
                best_distance = distance
                best_match_id = student_id
        
        # Check if best match is within tolerance
        if best_distance < tolerance:
            confidence = (1 - best_distance) * 100
            return True, best_match_id, confidence, "Face matched successfully!"
        else:
            return False, None, 0, "No matching face found in database"
            
    except Exception as e:
        return False, None, 0, f"Error during recognition: {str(e)}"


def get_image_path(student_id):
    """Get path to student's image"""
    return os.path.join(KNOWN_USER_DIR, f"{student_id}.jpg")


def image_exists(student_id):
    """Check if student's image exists"""
    return os.path.exists(get_image_path(student_id))


def delete_image(student_id):
    """Delete student's image and encoding"""
    image_path = os.path.join(KNOWN_USER_DIR, f"{student_id}.jpg")
    encoding_path = os.path.join(ENCODINGS_DIR, f"{student_id}.pkl")
    
    if os.path.exists(image_path):
        os.remove(image_path)
    if os.path.exists(encoding_path):
        os.remove(encoding_path)


def get_all_images():
    """Get list of all student IDs with images"""
    return [f.replace('.jpg', '') for f in os.listdir(KNOWN_USER_DIR) if f.endswith('.jpg')]