#best restoration model

import gradio as gr
import numpy as np
from ultralytics import YOLO
import cv2

model = YOLO("best.pt")

with gr.Blocks() as demo:
    with gr.Row():
      gr.Markdown("## Web app that segement the roof from house image using YLOV8")
    with gr.Row():
      gr.Markdown("### Please upload an image of the house, then click on Segement .")
    with gr.Row():  
      gr.Markdown("Then the yolov8 will segement the roof.")


    with gr.Row():
        input_image = gr.Image(source='upload', elem_id="input_image_upload", label="UploadImage")
        output_image = gr.Image(label="output")
    with gr.Row():
        segement_an = gr.Button("Segement")
    def segement_(input_):
        
        img = model.predict(source=input_,
                         stream=True, retina_masks=True)
        for result in img:
            mask = result.masks.cpu().numpy()
            masks = mask.masks.astype(bool)
            new = np.ones_like(input_, dtype=np.uint8)
            for m in masks:
                new[m] = input_[m]
            return new

    segement_an.click(segement_, inputs=[input_image], outputs=[output_image])
    demo.launch()