import gradio as gr
from ultralytics import YOLO

# Load the YOLO model
model = YOLO('best1.pt')

def predict(img, confidence_threshold):
    # Perform inference
    results = model(img)
    
    # Filter predictions based on the confidence threshold
    # The results[0].boxes.data contains the detection results, including confidence scores
    filtered_boxes = [box for box in results[0].boxes.data if box[4] >= confidence_threshold]
    
    # Plot the results (with the filtered detections)
    annotated_frame = results[0].plot(labels=filtered_boxes)
    
    return annotated_frame

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(label="Input Image", type="filepath"),
        gr.Slider(minimum=0, maximum=1, value=0.5, label="Confidence Threshold", step=0.01)
    ],
    outputs="image",
    title="Coin Detector",
    description="Upload an image to detect coins. Adjust the confidence threshold to filter results."
)

# Launch the Gradio interface
iface.launch(share=True)