import streamlit as st
from PIL import Image
from ultralytics import YOLO

# Load the YOLO model
@st.cache_resource
def load_model():
    model = YOLO('best (5).pt')  # Use the path to your trained model
    return model

# Prediction function
def predict(model, image):
    results = model(image)
    return results

# Load and display the header image
header_image = Image.open('historicalPIC.png')
st.image(header_image, use_column_width=True)

# Streamlit UI
st.title("Historical Places")

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

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    
    # Convert the image to RGB format if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Load the model
    model = load_model()
    
    # Run YOLO prediction
    results = predict(model, image)

    # Access class names
    class_names = model.names

    # Process and display results
    detected_objects = []
    for result in results:
        # Iterate over each detected object
        for bbox in result.boxes:
            x1, y1, x2, y2 = bbox.xyxy[0].tolist()
            conf = bbox.conf.item()
            cls = int(bbox.cls.item())
            detected_objects.append(f"Detected {class_names[cls]}")
    
    # Display detected objects' names and confidence scores
    st.subheader("Detection Results")
    if detected_objects:
        for obj in detected_objects:
            st.write(obj)
    else:
        st.write("No objects detected.")

    # Render the image with bounding boxes
    if results:
        try:
            results.render()  # Modify the image with bounding boxes
            img_with_boxes = Image.fromarray(results.imgs[0])
            st.image(img_with_boxes, caption='Detected Objects', use_column_width=True)
        except Exception:
            pass  # Ignore any errors in rendering