Spaces:
Sleeping
Sleeping
| 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 | |
| model = load_model("cpu") # Load the model on CPU by default | |
| 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() | |
| global model | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}") # Debug print | |
| model = model.to(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}", [] | |