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

# Initialize the YOLO model
model_path = 'yolov11x1.1-trained.pt'  # Ensure this model file is in the same directory
model = YOLO(model_path)

# Temporary fix: add placeholder names for missing classes
expected_classes = 13  # Set this to the correct number of classes
for i in range(expected_classes):
    if i not in model.names:
        model.names[i] = f"class_{i}"

def annotate_image(input_image_path, output_image_path, confidence=0.25):
    """Loads an image, runs YOLO model to detect skin issues with specified confidence, and saves annotated image."""
    
    # Load the image
    img = cv2.imread(input_image_path)
    if img is None:
        raise ValueError(f"Image at path {input_image_path} could not be loaded.")
    
    # Run YOLO model inference with the specified confidence threshold
    results = model.predict(img, conf=confidence)
    
    # Get the annotated image from results
    annotated_img = results[0].plot()
    
    # Save the annotated image to the specified output path
    cv2.imwrite(output_image_path, annotated_img)
    print(f"Annotated image saved at {output_image_path} with confidence threshold {confidence}")

# Streamlit UI
st.title("Skin Issue Detection with YOLO")
st.write("Upload an image to detect and annotate skin issues with a confidence threshold.")

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

# Confidence slider
confidence_threshold = st.slider("Confidence Threshold", min_value=0.0, max_value=1.0, value=0.25)

if uploaded_file is not None:
    # Save the uploaded file locally as 'test1.jpeg'
    input_image_path = 'test1.jpeg'
    with open(input_image_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    
    output_image_path = 'annotated_test1.jpeg'

    # Annotate image using your existing code function
    try:
        annotate_image(input_image_path, output_image_path, confidence=confidence_threshold)

        # Display the original and annotated images side by side
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Original Image")
            st.image(uploaded_file, use_column_width=True)
        
        with col2:
            st.subheader("Annotated Image")
            annotated_img = Image.open(output_image_path)
            st.image(annotated_img, use_column_width=True)
        
        # Provide a download link for the annotated image
        with open(output_image_path, "rb") as file:
            btn = st.download_button(
                label="Download Annotated Image",
                data=file,
                file_name="annotated_test1.jpeg",
                mime="image/jpeg"
            )
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")