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

# Define a class for YOLO-based crack detection
class YOLOPredictor:
    def __init__(self, model_path):
        # Initialize the YOLO model with the given model path
        self.model = YOLO(model_path)
        # Define class labels
        self.class_labels = [
            'alligator_crack_high', 'alligator_crack_low', 'alligator_crack_medium',
            'long_transverse_crack_high', 'long_transverse_crack_low', 'long_transverse_crack_medium'
        ]

    def predict(self, image):
        # Use the model to predict cracks in the image
        results = self.model.predict(image)
        
        # Debugging: Print results object to inspect its structure
        print("Results:", results)
        
        # Plot the results on the image
        result_img = results[0].plot()
        # Convert the image from BGR to RGB (OpenCV uses BGR by default)
        #result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
        
        # Extract class labels and confidence scores from the results
        confidences = []
        classes = []
        for box in results[0].boxes:
            confidences.append(box.conf.item())
            classes.append(int(box.cls.item()))  # Ensure class index is integer
        
        # Map class indices to class labels
        class_labels_mapped = [self.class_labels[idx] for idx in classes]
        
        return result_img, class_labels_mapped, confidences

def main():
    st.title("Thermal Crack Detection Dashboard")
    st.write("Upload an image to detect cracks and visualize the results.")

    # File uploader for the user to upload an image
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    # Initialize predictor outside the if block
    predictor = None

    # Check if a file has been uploaded
    if uploaded_file is not None:
        # Open the uploaded image
        image = Image.open(uploaded_file)
        # Convert the image to a NumPy array
        image = np.array(image)
        # Display the uploaded image
        st.image(image, caption='Uploaded Image', use_column_width=True)

        st.write("Running prediction...")
        # Create a YOLOPredictor object with the specified model path
        predictor = YOLOPredictor('model/best.pt')
        # Predict cracks in the uploaded image
        result_img, classes, confidences = predictor.predict(image)

        # Display the predicted image
        st.image(result_img, caption='Predicted Image', use_column_width=True)

        # Display detection classes and confidence scores
        if confidences and classes:
            st.write("Detection Classes and Confidence Scores:")
            for i, (cls, confidence) in enumerate(zip(classes, confidences)):
                st.write(f"Crack {i+1} - {cls}, Confidence: {confidence:.2f}")
        else:
            st.write("No detection classes and confidence scores available.")

        # Convert the result image to a format suitable for download
        pil_img = Image.fromarray(result_img)
        buf = io.BytesIO()
        pil_img.save(buf, format="PNG")
        byte_im = buf.getvalue()

        # Add a download button for the predicted image
        st.download_button(
            label="Download Annotated Image",
            data=byte_im,
            file_name="predicted_image.png",
            mime="image/png"
        )
    # Model information section
    st.sidebar.write("### Model Information")
    st.sidebar.write("Model: YOLO")
    st.sidebar.write("Version: v8")
    st.sidebar.write("Accuracy: 95%")
    
    if predictor is not None:
        st.sidebar.write(f"Total Classes: {len(predictor.class_labels)}")
        st.sidebar.write("### Class Labels:")
        for idx, label in enumerate(predictor.class_labels):
            st.sidebar.write(f"{idx}: {label}")
    else:
        st.sidebar.write("Please upload an image to display model information.")


if __name__ == "__main__":
    main()