import gradio as gr
from ultralytics import YOLO

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

    model = YOLO("model.pt")

    results = model.predict(source=img, save=False, show_labels=False, show_conf=False)
    
    # Handle both single result and list of results
    if not isinstance(results, (list, tuple)):
        results = [results]
    
    count = 0
    annot_img = None
    for i in results:
        count = len(i.boxes)
        annot_img = i.plot(labels=False, masks=False)
        # Convert BGR to RGB if needed
        if annot_img.shape[-1] == 3:  # Check if image has 3 channels
            annot_img = annot_img[..., ::-1]  # Reverse the color channels
    
    return str(count), annot_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="numpy")
        button.click(fn=predict, inputs=img, outputs=[data_output, img_output])

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