import streamlit as st
import PIL.Image as Image
import numpy as np
import pandas as pd
import requests
from io import BytesIO
from fastai.vision.all import load_learner  


options = st.selectbox("which model you wanna choose?", ("wbc classifier", "Blood Cell Detection with YOLOv8"))

if options == "wbc classifier":
    # Initialize Streamlit app
    st.title("White Blood Cell Classifier")
    
    # Add a description or subtitle
    st.markdown("""
    This app allows you to classify white blood cells from an uploaded image. 
    You can upload an image of a blood sample, and the app will predict the type of white blood cell present. 
    Choose from various cell types like eosinophil, lymphocyte, monocyte, and neutrophil.
    Note: To get the best results, please make sure there is only one WBC in the image. This model has not been trained on basophils.
    """)
    
    # Load the FastAI model for WBC identification
    fastai_model = load_learner('model1.pkl')
    
    # File uploader for image input
    uploaded_file = st.file_uploader("Upload an image for classification", type=["jpg", "png"])
    
    if uploaded_file:
        # Open the uploaded image
        image = Image.open(uploaded_file).convert('RGB')
    
        
        # Display the uploaded image with a caption
        st.image(image, caption="Reduced Size Image", use_column_width=False, width=150)  # 150 pixels wide
    
        # Perform inference with the FastAI model
        pred, idx, probs = fastai_model.predict(image)
    
        # Display a title for the results section
        st.subheader("White Blood Cell Classification Results")
    
        # Define categories for classification
        categories = ('EOSINOPHIL', 'LYMPHOCYTE', 'MONOCYTE', 'NEUTROPHIL')
    
        # Create a DataFrame with classification probabilities
        results_df = pd.DataFrame(
            {'Cell Type': categories, 'Probability': probs.tolist()}
        )
         # Highlight the most likely class
        most_likely_class = categories[idx]
        st.success(f"Predicted Class: {most_likely_class}")
    
        # Additional information about the probabilities
        st.write("Detailed Classification Results:")
        st.table(results_df)
    
        # Display the probabilities as a bar chart
        st.bar_chart(results_df.set_index('Cell Type'))
    
       
    else:
        st.warning("Upload an image to start classification.")

if options == "Blood Cell Detection with YOLOv8":
    # Initialize Streamlit app
    st.title("Blood Cell Detection with YOLOv8")
    
    # Load YOLO model
    model = YOLO('keremberke/yolov8m-blood-cell-detection')
    
    # Set model parameters
    model.overrides['conf'] = 0.25  # NMS confidence threshold
    model.overrides['iou'] = 0.45   # NMS IoU threshold
    model.overrides['agnostic_nms'] = False  # NMS class-agnostic
    model.overrides['max_det'] = 1000  # Maximum number of detections per image
    
    # Load the FastAI model for WBC identification
    fastai_model = load_learner('model1.pkl')
    
    # File uploader for image input
    uploaded_file = st.file_uploader("Upload an image for detection", type=["jpg", "png"])
    
    if uploaded_file:
        # Open the uploaded image
        image = Image.open(uploaded_file)
        
        # Perform inference
        results = model.predict(np.array(image))
        
        # Display results
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Render detection results
        rendered_image = render_result(model=model, image=image, result=results[0])
        
        # Show the rendered result
        st.image(rendered_image, caption="Detection Results", use_column_width=True)
        
        # Count the number of each cell type
        cell_counts = {"RBC": 0, "WBC": 0, "Platelets": 0}
        
        # Count cells and check for WBC
        has_wbc = False
        # Display details of detected boxes
        st.write("Detection Results:")
        for box in results[0].boxes:
            class_index = int(box.cls)  # Get the class index
            if class_index == 1:  # RBC
                cell_counts["RBC"] += 1
            elif class_index == 2:  # WBC
                cell_counts["WBC"] += 1
                has_wbc = True  # WBC detected
            elif class_index == 0:  # Platelets
                cell_counts["Platelets"] += 1
            
            # Display bounding box information
            #st.write(f"Bounding box: {box.xyxy}")
            #st.write(f"Confidence: {box.conf}")
            #st.write(f"Class: {box.cls}")
    
        # Display the counts of each cell type
        st.write("Cell Type Counts:")
        st.write(pd.DataFrame.from_dict(cell_counts, orient='index', columns=['Count']))
        # If a WBC is detected, run the second model
        if has_wbc:
            # Perform inference with the FastAI model
            pred, idx, probs = fastai_model.predict(image)
            st.write("White Blood Cell Classification:")
            categories = ('EOSINOPHIL', 'LYMPHOCYTE', 'MONOCYTE', 'NEUTROPHIL')
            results_dict = dict(zip(categories, map(float, probs)))
            st.write(results_dict)
    else:
        st.write("Upload an image to start detection.")