#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()