age_estimation/age_estimation.py

import torch

from .model import load_model
from utils.image_utils import preprocess_image, get_image_from_input
from utils.face_detector import load_face_detector
from .predict import predict_age


def age_estimation(input_type, uploaded_image, image_url, base64_string):
    """
    Estimates the age from an image input via file, URL, or base64 string.

    Args:
        input_type (str): The selected input method ("Upload File", "Enter URL", "Enter Base64").
        uploaded_image (PIL.Image.Image): The uploaded image (if input_type is "Upload File").
        image_url (str): The image URL (if input_type is "Enter URL").
        base64_string (str): The image base64 string (if input_type is "Enter Base64").

    Returns:
        tuple: A tuple containing:
            - str: A summary string of the estimated ages, or an error message.
            - list: A list of dictionaries, where each dictionary represents the age
                    estimation data for a detected face, or an empty list if no faces
                    were detected or an error occurred.
    """
    # Use the centralized function to get the image
    image = get_image_from_input(input_type, uploaded_image, image_url, base64_string)

    if image is None:
        print("Image is None after loading/selection for age estimation.")
        return "Error: Image processing failed or no valid input provided.", []

    try:
        face_detector = load_face_detector()

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = load_model(device)

        # Preprocess the image (convert PIL to numpy, ensure RGB)
        processed_image = preprocess_image(image)

        # Call predict_age with the processed image (NumPy array)
        age_data = predict_age(processed_image, model, face_detector, device)

        if age_data:
            # Create a summary string of all estimated ages
            age_summary = "Estimated Ages: " + ", ".join(
                [str(face["age"]) for face in age_data]
            )
            return age_summary, age_data
        else:
            return "No faces detected", []
    except Exception as e:
        print(f"Error in age estimation: {e}")
        return f"Error in age estimation: {e}", []