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

model = YOLO('best.pt')
mask_annotator = sv.MaskAnnotator()
box_annontator = sv.BoxAnnotator()


st.title("Yolov8 Solar Panel Segmentation")

uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"], accept_multiple_files=False)

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert('RGB')
    image_arr = np.array(image)
    
    result = model.predict(image)[0] # Use `image` because of RGB to BGR conversion 
    
    detections = sv.Detections.from_ultralytics(result)
    
    labels = [
        f"solar panel {confidence:0.2f}"
        for _, _, confidence, _, _
        in detections
    ]
    annotated_image = mask_annotator.annotate(image_arr, detections)
    annotated_image = box_annontator.annotate(annotated_image, detections, labels=labels)
    
    
    st.image(annotated_image, caption='Uploaded Image.', use_column_width=True)