import gradio as gr
from ultralytics import YOLO
import random
import os
from pathlib import Path

model_path = 'https://huggingface.co/mayrajeo/yolov8-deadwood/resolve/main/models/'

def run_models(
        im:gr.Image=None,
        model_type:gr.Dropdown='YOLOv8n',
        conf_thr:gr.Slider=0.25
):
    
    hp_model = YOLO(f'{model_path}{model_type}_hp.pt')
    hp_model.to(device='cpu')
    hp_result = hp_model(im[:,:,::-1], conf=conf_thr)
    hp_im = hp_result[0].plot()

    spk_model = YOLO(f'{model_path}{model_type}_spk.pt')
    spk_model.to(device='cpu')
    spk_result = spk_model(im[:,:,::-1], conf=conf_thr)
    spk_im = spk_result[0].plot()

    both_model = YOLO(f'{model_path}{model_type}_both.pt')
    both_model.to(device='cpu')
    both_result = both_model(im[:,:,::-1], conf=conf_thr)
    both_im = both_result[0].plot()

    return [
        (hp_im[:,:,::-1], 'HP'),
        (spk_im[:,:,::-1], 'SPK'),
        (both_im[:,:,::-1], 'HP+SPK')
    ]

ex_dir = Path('examples')

loc = gr.Textbox(label='Location')

desc_str = """
Demo application for YOLOv8 models for deadwood segmentation from RGB UAV imagery. Results are shown on three different models: HP is trained only with data from Hiidenportti, 
SPK only with data from Sudenpesänkangas and HP+SPK is trained with both sites.
"""

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown(desc_str)
    with gr.Row():
        with gr.Column(2):
            inp = gr.Image(label='Input image', sources='upload')
        with gr.Column(1):
            ex_list = random.sample([[ex_dir/i, i.split('_')[0]] for i in os.listdir(ex_dir)], 15)
            ex = gr.Examples(ex_list, inputs=[inp, loc],
                             cache_examples=False, examples_per_page=5,
                             label='Example UAV images')
        with gr.Column(1):
            loc.render()
            model = gr.Dropdown([
                        'YOLOv8n', 
                        'YOLOv8s', 
                        'YOLOv8m', 
                        'YOLOv8l', 
                        'YOLOv8x'
                        ],
                        value='YOLOv8n', label='Model')
            conf = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label='Confidence Threshold')
            btn = gr.Button()
    with gr.Row():
        with gr.Column():
            gallery = gr.Gallery(
                label='Predictions', show_label=True, elem_id='gallery',
                columns=[3], rows=[1], object_fit='contain', interactive=False
            ) 
    btn.click(run_models, [inp, model, conf], gallery)

if __name__ == '__main__': demo.launch(share=False)