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

# Clear GPU memory
torch.cuda.empty_cache()

# Load YOLO model
@st.cache_resource
def load_model():
    # Dynamically select device (GPU if available, otherwise CPU)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = YOLO('./best .pt').to(device)  # Replace with your model path
    return model

model = load_model()

# Function to make predictions and draw bounding boxes
def predict_and_draw(image):
    # Convert PIL image to OpenCV format
    img = np.array(image)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    # Resize image to YOLO input size (640x640)
    img_resized = cv2.resize(img, (640, 640))

    # Perform prediction
    results = model(img_resized)

    # Access the first result
    result = results[0]
    boxes = result.boxes  # Bounding boxes
    class_names = model.names  # Class names from the model
    img_with_boxes = img_resized.copy()  # Copy to draw on
    defect_list = []  # To store detected defect types and names

    # Draw bounding boxes and labels on the image
    for box in boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())  # Get bounding box coordinates
        conf = box.conf[0].item()  # Confidence score
        cls = int(box.cls[0].item())  # Class index
        label = f"{class_names[cls]} ({conf:.2f})"  # Class label with confidence

        defect_list.append(f"{class_names[cls]} - Confidence: {conf:.2f}")  # Add to list

        # Draw rectangle and label
        cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            img_with_boxes,
            label,
            (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 0, 0),
            2,
        )

    # Convert back to PIL for Streamlit display
    img_with_boxes = cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB)
    return Image.fromarray(img_with_boxes), defect_list

# Streamlit app
st.title("🚧 Road Defect Detection App 🚧")
st.markdown("Upload an image of a road to detect defects such as cracks, potholes, etc., with bounding boxes.")

# File uploader with a friendly description
uploaded_file = st.file_uploader("Upload an Image (JPG/PNG)", type=["jpg", "jpeg", "png"])

if uploaded_file:
    # Display input image and progress bar for prediction
    col1, col2 = st.columns([1, 1])  # Equal width for input and output columns

    with col1:
        st.subheader("Uploaded Image")
        st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)

    # Show a button for detection
    if st.button("Detect Defects"):
        with st.spinner("Detecting defects... Please wait."):
            # Show progress bar
            progress_bar = st.progress(0)
            result_image, defect_list = predict_and_draw(Image.open(uploaded_file))

            # Update progress
            progress_bar.progress(100)

            # Display result image with bounding boxes
            with col2:
                st.subheader("Detected Defects")
                st.image(result_image, caption="Detected Defects", use_container_width=True)

            # Display detected defects with confidence scores
            st.subheader("Detected Defects Details:")
            if defect_list:
                for defect in defect_list:
                    st.write(f"- {defect}")
            else:
                st.write("No defects detected.")

    else:
        st.warning("Click on 'Detect Defects' to analyze the image.")
else:
    st.info("Please upload an image to begin detection.")

# Add some footer information
st.markdown("""
    ---  
    🛠️ This app helps detect road defects using YOLO model.  
    📩 For feedback, contact us at: vaman2425@gmail.com
""")