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

# Load the YOLO model
@st.cache_resource
def load_model():
    model = YOLO("mosaic_medium_100_tiny_object.pt")  
    model.to('cpu')  
    return model

model = load_model()

def process_image(image):
    # Convert PIL Image to numpy array
    image_np = np.array(image)
    
    # Convert RGB to BGR (OpenCV uses BGR)
    image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    def callback(image_slice: np.ndarray) -> sv.Detections:
        result = model(image_slice)[0]
        return sv.Detections.from_ultralytics(result)

    slicer = sv.InferenceSlicer(callback=callback, slice_wh=(256, 256), iou_threshold=0.8)

    detections = slicer(image_cv2)

    # Filter detections for building class (assuming class_id 2 is for buildings)
    building_detections = detections[detections.class_id == 2]

    label_annotator = sv.LabelAnnotator()
    box_annotator = sv.BoxAnnotator()

    annotated_image = box_annotator.annotate(scene=image_cv2.copy(), detections=building_detections)
    annotated_image = label_annotator.annotate(scene=annotated_image, detections=building_detections)

    # Convert BGR back to RGB for displaying in Streamlit
    return cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

def main():
    st.title("Building Detection App")

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

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        if st.button("Detect Buildings"):
            with st.spinner("Processing..."):
                result_image = process_image(image)
                st.image(result_image, caption="Processed Image", use_column_width=True)

if __name__ == "__main__":
    main()