import streamlit as st
from ultralytics import YOLO
from PIL import Image
import numpy as np

# Khởi tạo model YOLOv8s
model = YOLO("yolov10n.pt")  # Đường dẫn tới model YOLOv10n đã huấn luyện sẵn

st.title("YOLOv10 Image Prediction")

# Tải tệp ảnh đầu vào
uploaded_file = st.file_uploader("Chọn một tệp ảnh", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Hiển thị ảnh đầu vào
    image = Image.open(uploaded_file)
    st.image(image, caption="Ảnh đầu vào", use_column_width=True)
    # Nút "Predict" để chạy mô hình
    if st.button("Predict"):
        # Chạy dự đoán với YOLO
        results = model.predict(source=np.array(image))
        # Lấy ảnh kết quả từ mô hình
        result_image = results[0].plot()  # Vẽ bounding box lên ảnh gốc
        # Hiển thị ảnh kết quả
        st.image(result_image, caption="Kết quả dự đoán", use_column_width=True)