age_estimation/age_estimation.py

import os
import tempfile
import torch
import dlib

from PIL import Image

from .model import load_model
from utils.image_utils import load_image, 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:
        str: The estimated age, or an error message.
    """
    # 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:
            # Assuming age_data is a list of dictionaries, and we take the first face's age
            return f"Estimated Age: {age_data[0]['age']}"
        else:
            return "No faces detected"
    except Exception as e:
        print(f"Error in age estimation: {e}")
        return f"Error in age estimation: {e}"