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

# Set page config
st.set_page_config(page_title="Building Detection App", page_icon="🏢", layout="wide")

# Custom CSS with theme compatibility
st.markdown("""
<style>
    .reportview-container {
        background: var(--background-color);
    }
    .main {
        background-color: var(--background-color);
        padding: 2rem;
        border-radius: 10px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        font-weight: bold;
        border: none;
        border-radius: 5px;
        padding: 0.5rem 1rem;
        transition: all 0.3s;
    }
    .stButton>button:hover {
        background-color: #45a049;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
    }
    .upload-box {
        border: 2px dashed #4CAF50;
        border-radius: 10px;
        padding: 2rem;
        text-align: center;
    }
    .theme-text {
        color: var(--text-color);
    }
</style>
""", unsafe_allow_html=True)

# Load the YOLO model
@st.cache_resource
def load_model():
    model = YOLO("mosaic_medium_100_tiny_object.pt")  # Update this to the filename of your model
    model.to('cpu')  # Ensure the model is on 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")
    st.markdown('<p class="theme-text">Upload an image to detect buildings using our advanced AI model.</p>', unsafe_allow_html=True)

    col1, col2 = st.columns(2)

    with col1:
        st.markdown('<h3 class="theme-text">Upload Image</h3>', unsafe_allow_html=True)
        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)
                    with col2:
                        st.markdown('<h3 class="theme-text">Detection Results</h3>', unsafe_allow_html=True)
                        st.image(result_image, caption="Processed Image", use_column_width=True)
                    st.success("Detection completed successfully!")
        else:
            st.markdown(
                """
                <div class="upload-box theme-text">
                    <h3>📁 Upload an image to get started</h3>
                    <p>Supported formats: JPG, JPEG, PNG</p>
                </div>
                """,
                unsafe_allow_html=True
            )

    with col2:
        if uploaded_file is None:
            st.markdown('<h3 class="theme-text">How it works</h3>', unsafe_allow_html=True)
            st.markdown(
                """
                <p class="theme-text">
                1. <strong>Upload</strong> an image using the file uploader on the left.<br>
                2. Click the <strong>Detect Buildings</strong> button to process the image.<br>
                3. View the results with bounding boxes around detected buildings.<br><br>
                Our AI model is trained to identify various types of buildings in different environments.
                </p>
                """,
                unsafe_allow_html=True
            )

    st.markdown("---")
    #st.markdown('<p class="theme-text"></p>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()