import streamlit as st
import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.vgg16 import preprocess_input
from tensorflow.keras.preprocessing import image
from werkzeug.utils import secure_filename
import os

# Loading Models
braintumor_model = load_model('models/brain_tumor_binary.h5')

# Configuring Streamlit
UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
st.set_page_config(page_title="Brain Tumor Prediction App", page_icon=":brain:")

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def preprocess_imgs(set_name, img_size):
    set_new = []
    for img in set_name:
        img = cv2.resize(img, dsize=img_size, interpolation=cv2.INTER_CUBIC)
        set_new.append(preprocess_input(img))
    return np.array(set_new)

def crop_imgs(set_name, add_pixels_value=0):
    set_new = []
    for img in set_name:
        gray = cv2.GaussianBlur(img, (5, 5), 0)
        thresh = cv2.threshold(gray, 45, 255, cv2.THRESH_BINARY)[1]
        thresh = cv2.erode(thresh, None, iterations=2)
        thresh = cv2.dilate(thresh, None, iterations=2)
        cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cnts = cnts[0] if len(cnts) == 2 else cnts[1]
        c = max(cnts, key=cv2.contourArea)
        extLeft = tuple(c[c[:, :, 0].argmin()][0])
        extRight = tuple(c[c[:, :, 0].argmax()][0])
        extTop = tuple(c[c[:, :, 1].argmin()][0])
        extBot = tuple(c[c[:, :, 1].argmax()][0])
        ADD_PIXELS = add_pixels_value
        new_img = img[extTop[1] - ADD_PIXELS:extBot[1] + ADD_PIXELS,
                      extLeft[0] - ADD_PIXELS:extRight[0] + ADD_PIXELS].copy()
        set_new.append(new_img)
    return np.array(set_new)

# Function to preprocess the image
def preprocess_image(img_array):
    img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
    img = image.array_to_img(img_array, target_size=(200, 200))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array /= 255.0  # Normalize the image
    return img_array

def predict_braintumor(img_input):
    if isinstance(img_input, str):  # File path (for Streamlit)
        img_path = img_input
        img_gray = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    else:  # NumPy array (for Gradio)
        img_array = img_input.astype(np.uint8)
        img_gray = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)

    # Crop and preprocess the grayscale image
    img_processed = crop_imgs([img_gray])
    img_processed = preprocess_imgs(img_processed, (224, 224))

    # Make prediction
    pred = braintumor_model.predict(img_processed)

    # Handle binary decision
    pred = 1 if pred[0] >= 0.5 else 0

    return pred

def main():
    st.title("Brain Tumor Prediction App")

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

    if uploaded_file is not None:
        st.image(uploaded_file, caption="Uploaded Image.", use_column_width=True)
        st.write("")
        st.write("Classifying...")

        # Read the contents of the uploaded file or the input image
        file_contents = uploaded_file.read()
        uploaded_array = np.asarray(bytearray(file_contents), dtype=np.uint8)
        uploaded_img = cv2.imdecode(uploaded_array, -1)

        # Save the uploaded file
        if isinstance(file_contents, bytes):
            filename = secure_filename(uploaded_file.name)
            file_path = os.path.join(UPLOAD_FOLDER, filename)

            with open(file_path, "wb") as f:
                f.write(file_contents)

            # Make prediction
            result = predict_braintumor(file_path)
        else:
            # Make prediction for Gradio (direct input of image)
            result = predict_braintumor(uploaded_img)

        # Handle binary decision
        result_text = "Brain Tumor Found!" if result == 1 else "Brain Tumor Not Found!"
        
        # Display prediction
        st.subheader("Prediction:")
        st.success(result_text)

if __name__ == "__main__":
    # Streamlit part
    if "streamlit" in st.__version__.lower():
        main()

    # Gradio part
    iface = gr.Interface(
        fn=predict_braintumor,
        inputs="image",
        outputs="text",
        examples=[["examples/1 no.jpg"]],
        enable_queue=True
    )
    iface.launch()