import streamlit as st
from ultralytics.utils.plotting import Annotator
import numpy as np
import cv2
from PIL import Image
import io
import yolov9

# Load the YOLOv9 model
model = yolov9.load('fracture.pt', device="cpu")
model.conf = 0.1  # Lowered the confidence threshold
model.iou = 0.45

def Predict(img):
    # Convert the image to an OpenCV format
    img_array = np.array(img)
    img = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

    # Perform prediction
    results = model(img, size=640)
    print(results)
    annotator = Annotator(img, line_width=2, example=str('Fracture'))

    for result in results.xyxy[0]:
        xmin, ymin, xmax, ymax, confidence, class_id = result
        label = results.names[int(class_id)]
        confidence = float(confidence)

        # Annotate the image
        annotator.box_label([xmin, ymin, xmax, ymax], f"{label} {confidence:.2f}", color=(255, 0, 0))

    annotated_img = annotator.result()
    return annotated_img

# Streamlit app
st.title("YOLOv9 Fracture Detection")

# Image upload section
uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    # Convert the uploaded file to a JPEG image
    image = Image.open(uploaded_file).convert("RGB")
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format="JPEG")
    image = Image.open(io.BytesIO(img_byte_arr.getvalue()))

    # Display the uploaded image
    st.image(image, caption='Uploaded Image.', use_column_width=True)

    # Run the prediction
    st.write("Detecting fractures...")
    annotated_image = Predict(image)

    # Convert the annotated image back to RGB for display in Streamlit
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

    # Display the annotated image
    st.image(annotated_image, caption='Annotated Image.', use_column_width=True)