import streamlit as st
import cv2
import numpy as np
import tempfile
import os
from ultralytics import YOLO

# Load the YOLOv8 model
try:
    model = YOLO('best.pt')  # Ensure the path to your trained YOLOv8 model weights is correct
except FileNotFoundError:
    st.error("Model weights file 'best.pt' not found. Please ensure the model file is in the correct directory.")
    st.stop()

# Predefined colors for each class
COLORS = {
    "high-pneumonia": [255, 0, 0],    # Red
    "low-pneumonia": [0, 255, 0],     # Green
    "no-pneumonia": [0, 0, 255]       # Blue
}

# Define a constant image size
IMAGE_SIZE = (640, 640)  # Width, Height

def segment_image(image):
    results = model(image)
    return results

def display_segmented_image(image, results):
    detected_classes = set()
    if results and results[0].masks is not None and results[0].boxes is not None:
        masks = results[0].masks.data.cpu().numpy()
        boxes = results[0].boxes.data.cpu().numpy()
        class_ids = results[0].boxes.cls.cpu().numpy()

        for mask, box, class_id in zip(masks, boxes, class_ids):
            class_name = model.names[int(class_id)]
            detected_classes.add(class_name)

        # If low-pneumonia or high-pneumonia is detected, ignore no-pneumonia masks
        if 'low-pneumonia' in detected_classes or 'high-pneumonia' in detected_classes:
            detected_classes.discard('no-pneumonia')

        for mask, box, class_id in zip(masks, boxes, class_ids):
            class_name = model.names[int(class_id)]
            if class_name not in detected_classes:
                continue
            mask = mask.astype(np.uint8)
            mask_resized = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
            color = COLORS.get(class_name, [255, 255, 255])  # Default to white if class not found
            image[mask_resized == 1] = color
            x1, y1, x2, y2 = box[:4].astype(int)
            # Draw a rectangle behind the text for better visibility
            text_size = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2)[0]
            text_x = x1
            text_y = y1 - 10
            cv2.rectangle(image, (text_x, text_y - text_size[1]), (text_x + text_size[0], text_y), color, -1)
            cv2.putText(image, class_name, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
            # Draw the bounding box rectangle
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
    else:
        st.warning("No objects detected.")

    # Resize the final image to the constant size
    resized_image = cv2.resize(image, IMAGE_SIZE, interpolation=cv2.INTER_LINEAR)
    
    return resized_image, detected_classes


st.title('X-ray Segmentation Project')

# Add an author section in the sidebar
st.sidebar.title('About the Author')
st.sidebar.markdown("""
    **Author Name**: Makhammadjonov Izzatullokh 
    **Email**: izzatullokhm@gmail.com  
    <a href="https://github.com/Izzatullokh24" target="_blank"><i class="fab fa-github"></i> GitHub</a>  
    <a href="https://www.linkedin.com/in/izzatullokh-makhammadjonov-242042195/" target="_blank"><i class="fab fa-linkedin"></i> LinkedIn</a>  
    <style>
        .fab {
            font-size: 24px;
            margin-right: 10px;
        }
    </style>
    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/css/all.min.css">
    Izzatullokh is a machine learning engineer with a passion for computer vision and deep learning.
""", unsafe_allow_html=True)

uploaded_file = st.file_uploader("Upload an image or video", type=["jpg", "jpeg", "png", "mp4", "avi"])

if uploaded_file is not None:
    file_type = uploaded_file.type.split('/')[0]

    try:
        if file_type == 'image':
            file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
            image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
            results = segment_image(image)
            segmented_image, detected_classes = display_segmented_image(image, results)
            st.image(segmented_image, caption='Segmented Image', use_column_width=True)

             # Display pneumonia information
            if detected_classes:
                st.subheader('Diagnosis:')
                diagnosis_html = ""
                if 'high-pneumonia' in detected_classes:
                    diagnosis_html += '<p style="color: red; font-size: 20px;">The patient has <strong>high pneumonia</strong>.</p>'
                if 'low-pneumonia' in detected_classes:
                    diagnosis_html += '<p style="color: green; font-size: 20px;">The patient has <strong>low pneumonia</strong>.</p>'
                if 'no-pneumonia' in detected_classes:
                    diagnosis_html += '<p style="color: blue; font-size: 20px;">The patient does <strong>not have pneumonia</strong>.</p>'
                st.markdown(diagnosis_html, unsafe_allow_html=True)
            else:
                st.write('No pneumonia detected.')

        elif file_type == 'video':
            tfile = tempfile.NamedTemporaryFile(delete=False)
            tfile.write(uploaded_file.read())
            video_path = tfile.name
            segmented_frames = segment_video(video_path)
            stframe = st.empty()
            for frame, results in segmented_frames:
                segmented_frame, detected_classes = display_segmented_image(frame, results)
                stframe.image(segmented_frame, channels="BGR")
            os.remove(video_path)

        else:
            st.error("Unsupported file format. Please upload a jpg, jpeg, png, mp4, or avi file.")

    except Exception as e:
        st.error(f"An error occurred: {e}")

else:
    st.info("Please upload a file to proceed.")