import streamlit as st
from PIL import Image, ImageDraw, ImageFont
from transformers import pipeline

# Load YOLO pipeline
yolo_pipe = pipeline("object-detection", model="hustvl/yolos-small")

# Function to draw bounding boxes and labels on image
def draw_boxes(image, predictions):
    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()

    for pred in predictions:
        label = pred['label']
        score = pred['score']
        box = pred['box']
        xmin, ymin, xmax, ymax = box.values()
        draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=2)
        draw.text((xmin, ymin), f"{label} ({score:.2f})", fill="red", font=font)

    return image

def main():
    st.title("YOLO Object Detection App")

    # Instructions
    st.sidebar.markdown("""
        📸 **YOLO Object Detection App**

        - Check "Use Camera" to capture images live.
        - Upload an image or use the camera.
        - Detects objects using YOLO model (hustvl/yolos-small).
        - Displays annotated image with bounding boxes.
        - Enjoy object detection in real-time or with uploaded images.
    """)

    # Use camera input
    camera = st.sidebar.checkbox("Use Camera", value=False)

    # Upload image
    if camera:
        uploaded_image = st.sidebar.camera_input("Take a picture 📸")    
    else:
        uploaded_image = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_image is not None:
        # Display uploaded image
        st.image(uploaded_image, caption="Captured/Uploaded Image", use_column_width=True)
        st.success("Image displayed successfully.")

        # Process image with YOLO
        image = Image.open(uploaded_image)
        with st.spinner('Wait for it...'):
            predictions = yolo_pipe(image)
            st.success("YOLO running successfully.")

        # Draw bounding boxes and labels
        image_with_boxes = draw_boxes(image.copy(), predictions)
        st.success("Bounding boxes drawn.")

        # Display annotated image
        st.image(image_with_boxes, caption="Annotated Image", use_column_width=True)

if __name__ == "__main__":
    main()