from keras.models import load_model
import cv2
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import gradio as gr
import numpy as np
from yolo_model import Predict

my_model=load_model('Liver_model.h5',compile=True)
heart_model=load_model('Chicken_Heart_model.h5',compile=True)
lu_model=load_model('Lungs_model.h5',compile=True)
auth_model=load_model('update_postmortem_auth_model.h5',compile=True)
heart_class_name = {0: 'Dilation(eccentric)', 1: 'Hepatoma', 2: 'Hypertrophy(concentric)', 3: 'Hypertrophy(physiological)', 4: 'Infraction Damage', 5: 'Normal'}
heart_result = {0: 'Critical', 1: 'Critical', 2: 'Critical', 3: 'Critical', 4: 'Critical', 5: 'Normal'}
heart_recommend = {0: 'panadol', 1: 'peracetamol', 2: 'ponston', 3: 'brofon', 4: 'brofon', 5: 'No Need'}

def Heart_Disease_prediction(img):
    img = cv2.imread(img)
    img = img.reshape((1, img.shape[0], img.shape[1], img.shape[2]))

    # Create the data generator with desired properties
    datagen = ImageDataGenerator(
        rotation_range=30,
        width_shift_range=0.1,
        height_shift_range=0.1,
        shear_range=0.1,
        zoom_range=0.1,
        horizontal_flip=True,
        fill_mode="nearest",
    )
    # Generate a batch of augmented images (contains only the single image)
    augmented_images = datagen.flow(img, batch_size=1)
    # Get the first (and only) augmented image from the batch
    augmented_img = next(augmented_images)[0]
    img = cv2.resize(augmented_img.astype(np.uint8), (128, 128))
    class_no = heart_model.predict(img.reshape(1, 128, 128, 3)).argmax()
    name = "Heart Organ: " + heart_class_name.get(class_no)
    result = "Heart Organ Status: " + heart_result.get(class_no)
    recommend = "Heart Organ Recommendation: " + heart_recommend.get(class_no)
    return name, result, recommend

liver_class_num = {0: 'Healthy', 1: 'Un-Healthy'}
liver_result = {0: 'Normal', 1: 'Critical'}
liver_recommend = {0: 'No need Medicine', 1: 'Panadol'}

def Liver_Predict(image):
    image=cv2.imread(image)
    image = cv2.resize(image, (224, 224))
    class_no = my_model.predict(image.reshape(1, 224, 224, 3)).argmax()
    class_name = "Liver Organ: " + liver_class_num.get(class_no)
    liver_class_result = "Liver Organ Status: " + liver_result.get(class_no)
    liver_class_recommend = "Liver Organ Recommendation: " + liver_recommend.get(class_no)
    return class_name, liver_class_result, liver_class_recommend

lung_classes = {
    0: 'Lungs of infected chickens showing congestion, hemorrhage and consolidation with traces of fibrin at 24 hpi (hours post-infection)',
    1: 'gradual paleness and reduction in size of lungs at 2 dpi (days post-infection)',
    2: 'gradual paleness and reduction in size of lung at 3 dpi (days post-infection)',
    3: 'severe congestion, hemorrhage, and gradual shrinking of lungs at 4 dpi (days post-infection)',
    4: 'severe congestion, hemorrhage, and gradual shrinking of lungs at 5 dpi (days post-infection)'
}
lung_result = {0: 'critical', 1: 'critical', 2: 'critical', 3: 'critical', 4: 'critical'}
lung_recommend = {0: 'panadol', 1: 'peracetamol', 2: 'ponston', 3: 'brofon', 4: 'brofon'}

def Lungs_predict(image):
    image = cv2.resize(cv2.imread(image), (224, 224))
    lung_no = lu_model.predict(image.reshape(1, 224, 224, 3)).argmax()
    lung_disease_name = "Lung Organ: " + lung_classes.get(lung_no)
    lung_r = "Lung Organ Status: " + lung_result.get(lung_no)
    lung_re = "Lung Organ Recommendation: " + lung_recommend.get(lung_no)
    return lung_disease_name, lung_r, lung_re

def main(Image):
    liver_name,liver_r,liver_re,heart_n,heart_r,heart_re,lung_d,lung_r,lung_re='Liver Organ: Not Detected','Liver Organ:N/A','Liver Organ:N/A','Heart Organ: Not Detected','Heart Organ: N/A','Heart Organ: N/A','Lungs Organ: Not Detected','Lungs Organ: N/A','Lungs Organ: N/A'
    img = cv2.resize(Image, (224, 224))
    indx = auth_model.predict(img.reshape(1, 224, 224, 3)).argmax()
    if indx == 0:
        liver_name,liver_r,liver_re,heart_n,heart_r,heart_re,lung_d,lung_r,lung_re='Liver Organ: Not Detected','Liver Organ: N/A','Liver Organ: N/A','Heart Organ: Not Detected','Heart Organ: N/A','Heart Organ: N/A','Lungs Organ: Not Detected','Lungs Organ: N/A','Lungs Organ: N/A'
        return liver_name,liver_r,liver_re,heart_n,heart_r,heart_re,lung_d,lung_r,lung_re
    else:
        img_name_list, labels = Predict(Image)
        if len(labels) > 0:
            if labels[0] is not None:
                if labels[0]['label'] == 'Liver':
                    liver_name, liver_r, liver_re = Liver_Predict('Liver.jpg')
                if labels[0]['label'] == 'Heart':
                    heart_n, heart_r, heart_re = Heart_Disease_prediction(img_name_list[0])
                if labels[0]['label'] == 'Lung':
                    lung_d, lung_r, lung_re = Lungs_predict(img_name_list[0])
            if len(labels) > 1 and labels[1] is not None:
                if labels[1]['label'] == 'Liver':
                    liver_name, liver_r, liver_re = Liver_Predict(Image)
                if labels[1]['label'] == 'Heart':
                    heart_n, heart_r, heart_re = Heart_Disease_prediction(img_name_list[1])
                if labels[1]['label'] == 'Lung':
                    lung_d, lung_r, lung_re = Lungs_predict(img_name_list[1])
        return liver_name, liver_r, liver_re, heart_n, heart_r, heart_re, lung_d, lung_r, lung_re

interface = gr.Interface(fn=main, inputs='image', outputs=[
    gr.components.Textbox(label="Heart Disease Name"),
    gr.components.Textbox(label="Heart result Name"),
    gr.components.Textbox(label="Heart recommend"),
    gr.components.Textbox(label="Liver Disease Name"),
    gr.components.Textbox(label="Liver result Name"),
    gr.components.Textbox(label="Liver recommend"),
    gr.components.Textbox(label="Lung Disease Name"),
    gr.components.Textbox(label="Lung result Name"),
    gr.components.Textbox(label="Lung recommend")
], title="Postmortem")

interface.launch(debug=True)