import os
import gradio as gr
from age_estimator.mivolo.demo_copy import main as age_estimation_main

def process_age_estimation(video_file):
    # Validate the input file path
    if not video_file or not os.path.isfile(video_file):
        return {'error': 'Invalid video path'}

    # Define parameters for age estimation
    output_folder = 'output'
    detector_weights = 'age_estimator/mivolo/models/yolov8x_person_face.pt'
    checkpoint = 'age_estimator/mivolo/models/model_imdb_cross_person_4.22_99.46.pth.tar'
    device = 'cpu'
    with_persons = True
    disable_faces = False
    draw = True
    
    # Run age estimation
    absolute_age, lower_bound, upper_bound = age_estimation_main(
        video_file, output_folder, detector_weights, checkpoint, device, with_persons, disable_faces, draw
    )

    # Compile results
    results = {
        'Age Range': f"{lower_bound} - {upper_bound}"
    }

    return results

# Define Gradio interface for Age Estimation
gr.Interface(
    fn=process_age_estimation,
    inputs=gr.Video(label="Upload a video file"),
    outputs="json",
    title="Age Estimation"
).launch(server_name="0.0.0.0")