import os
import shutil
import gradio as gr
from ultralytics import YOLO

def predict(img):
    path = img.split("\\")[-1].split(".")[0]
    print("path", path)

    if os.path.exists(r".\runs\segment\predict"):
        shutil.rmtree(r".\runs\segment\predict")
    model = YOLO("model.pt")

    results = model.predict(source=img,save=True, show_labels=False, show_conf=False)
    count = 0
    for i in results:
        count = len(i.boxes)
    cur_path = os.getcwd()
    result_img = os.path.join(cur_path,r"runs\segment\predict",f"{path}.jpg")
    print("result img",result_img)
    return str(count), result_img

with gr.Blocks(title="Pill Counter") as demo:

    with gr.Row():
        with gr.Column():
            img = gr.Image(type="filepath",format="jpg", height=500, width=700)

            button = gr.Button()
        with gr.Column():
            data_output = gr.Textbox()
            img_output = gr.Image(type="filepath")
        button.click(fn=predict, inputs=img, outputs=[data_output ,img_output])

if __name__ == "__main__":
    demo.launch(allowed_paths=["runs"])