import streamlit as st
import degirum as dg
from PIL import Image
import degirum_tools

# hw_location: Where you want to run inference.
#     Use "@cloud" to use DeGirum cloud.
#     Use "@local" to run on local machine.
#     Use an IP address for AI server inference.
hw_location = "@cloud"

# face_model_zoo_url: URL/path for the face model zoo.
#     Use cloud_zoo_url for @cloud, @local, and AI server inference options.
#     Use '' for an AI server serving models from a local folder.
#     Use a path to a JSON file for a single model zoo in case of @local inference.
face_model_zoo_url = "https://cs.degirum.com/degirum/ultralytics_v6"

# face_model_name: Name of the model for face detection.
face_model_name = "yolov8n_relu6_face--640x640_quant_n2x_orca1_1"

# age_model_zoo_url: URL/path for the age model zoo.
age_model_zoo_url = "https://cs.degirum.com/degirum/openvino"

# age_model_name: Name of the model for age detection.
age_model_name = "yolov8n_relu6_age--256x256_float_openvino_cpu_1"

# Connect to AI inference engine getting token from env.ini file
face_zoo = dg.connect(hw_location, face_model_zoo_url, token=st.secrets["DG_TOKEN"])
age_zoo = dg.connect(hw_location, age_model_zoo_url, token=st.secrets["DG_TOKEN"])

# Load models
face_model = face_zoo.load_model(face_model_name, 
                                 image_backend='pil',
                                 overlay_color=(255,0,0),
                                 overlay_line_width=2,
                                 overlay_font_scale=1.5                                
                                )
age_model= age_zoo.load_model(age_model_name, 
                              overlay_show_probabilities=True, 
                              overlay_show_labels=True, 
                              image_backend='pil')
# Create a compound cropping model with 50% crop extent
crop_model = degirum_tools.CroppingAndClassifyingCompoundModel(
    face_model, age_model, 50.0
    )

st.title('DeGirum Cloud Platform Demo of Face Detection and Age Regression Models')

st.text('Upload an image. Then click on the submit button')
with st.form("model_form"):
    uploaded_file=st.file_uploader('input image')
    submitted = st.form_submit_button("Submit")
    if submitted:
        image = Image.open(uploaded_file)
        image.thumbnail((640,640), Image.Resampling.LANCZOS)
        inference_results=crop_model(image)
        st.image(inference_results.image_overlay,caption='Image with Bounding Boxes')