import gradio as gr
from deepforest import main
import cv2

def show_trees(img_path):
    try:
        model = main.deepforest()
        model.use_release()
    except Exception as model_error:
        raise gr.Error(f"Error initializing the deepforest model: {model_error}")

    img=cv2.imread(img_path)
    
    if img is None:
        raise gr.Error(f"Image path is not valid.")
    
    if img.shape[0]<1000 and img.shape[1]<1000:
        img = model.predict_image(path=img_path, return_plot=True)
    elif img.shape[0]>1000 or img.shape[1]>1000:
        img=model.predict_image(path=img_path,return_plot=True,thickness=2)
    else:
        img = model.predict_image(path=img_path, return_plot=True,thickness=3) 
    
    if img is None:
        raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
   
    return img

def show_birds(img_path):
    try:
        model = main.deepforest()
        model.use_bird_release()
    except Exception as model_error:
        raise gr.Error(f"Error initializing the deepforest model: {model_error}")
    img=cv2.imread(img_path)
    
    if img is None:
        raise gr.Error(f"Image path is not valid.")
    
    if img.shape[0]<1000 and img.shape[1]<1000:
        img = model.predict_image(path=img_path, return_plot=True)
    elif img.shape[0]>1000 or img.shape[1]>1000:
        img=model.predict_image(path=img_path,return_plot=True,thickness=2)
    else:
        img = model.predict_image(path=img_path, return_plot=True,thickness=3) 
    
    if img is None:
        raise gr.Error("No predictions were made. Check your test image. Ensure it conists")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
   
    return img

with gr.Blocks() as demo:
    gr.Markdown('# Deepforest')
    gr.Markdown('## Tree Detection Model')
    gr.Markdown('### Predict trees')
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="filepath")
        with gr.Column():
            output_image = gr.Image(label="Predicted image")

    submit_button_trees = gr.Button("Predict trees")
    submit_button_trees.click(show_trees, inputs=input_image, outputs=output_image)
    
    gr.Markdown('### Predict birds')
    with gr.Row():
        with gr.Column():
            input_image=gr.Image(label="Input image",type="filepath")
        with gr.Column():
            output_image=gr.Image(label="Predicted Image")
            
    submit_button_birds = gr.Button("Predict birds")
    submit_button_birds.click(show_birds,inputs=input_image,outputs=output_image)

if __name__ == "__main__":
    demo.launch()