import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cv2
import keras
import gradio as gr

SHAPE = (224, 224, 3)

predictor_disease_risk = keras.models.load_model('predictor_Disease_Risk.h5')
predictor_dr = keras.models.load_model('predictor_DR.h5')
predictor_mh = keras.models.load_model('predictor_MH.h5')
predictor_odc = keras.models.load_model('predictor_ODC.h5')
predictor_tsln = keras.models.load_model('predictor_TSLN.h5') 
predictor_dn = keras.models.load_model('predictor_DN.h5')
predictor_armd = keras.models.load_model('predictor_ARMD.h5')
predictor_mya = keras.models.load_model('predictor_MYA.h5')
predictor_brvo = keras.models.load_model('predictor_BRVO.h5')


def cut_and_resize(image):
  LOW_TOL = 20
  img_bw = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  img_bw[img_bw<=LOW_TOL] = 0
  y_nonzero, x_nonzero = np.nonzero(img_bw)
  image = image[np.min(y_nonzero):np.max(y_nonzero), np.min(x_nonzero): np.max(x_nonzero), ]
  return cv2.resize(image, SHAPE[:2], interpolation = cv2.INTER_LINEAR)
 
def simple_normalizer(X):
  return X / 255.0

def predict (image_path):
  image = simple_normalizer(cut_and_resize(cv2.imread(image_path)))
  result = predictor_disease_risk.predict(np.array([image]))[0][0] 
  
  dr = predictor_dr.predict(np.array([image]))[0][0]  
  mh = predictor_mh.predict(np.array([image]))[0][0]
  odc = predictor_odc.predict(np.array([image]))[0][0] 
  tsln = predictor_tsln.predict(np.array([image]))[0][0]
  dn = predictor_dn.predict(np.array([image]))[0][0]
  armd = predictor_armd.predict(np.array([image]))[0][0]
  mya = predictor_mya.predict(np.array([image]))[0][0]
  brvo = predictor_brvo.predict(np.array([image]))[0][0]
  
  diseases = {
      'DR'  : float(dr),
      'MH'  : float(mh),
      'ODC' : float(odc),
      'DN'  : float(dn),
      'TSLN': float(tsln),
      'ARMD': float(armd),
      'MYA' : float(mya),
      'BRVO': float(brvo)    
  }
  
  to_delete = []
  for k,v in diseases.items():
    if v < 0.05:
      to_delete.append(k)
  
  for k in to_delete:
    del diseases[k]
  
  if len(diseases) == 0:
    diseases = {'No specific disease': 0.0}
  
  
  return (
    {'Enferma': float(result), 'Sana': 1 - float(result)}, diseases
  )
 
title = 'Retinal Disease Predictor'
description = 'Modelo de deep learning que permite clasificar imágenes de la retina en patológicas y no patológicas. Si detecta una retina enferma, realiza un diagnóstico de la enfermedad concreta entre las siguientes: Diabetic Retinopathy (DR), Media Haze (MH), Optic Disk Cupping (ODC), Drusen (DN), Tessellation (TSLN), Age Related Macular Disease (ARMD), Myopia (MYA), Branch Retinal Vein Occlusion (BRVO)  . Las imágenes deben tener fondo negro.'
article = 'Proyecto HORUS (Helping Oftalmoscopy of Retina Using Supervised Learning'

interface = gr.Interface(
  predict,
  inputs = [gr.inputs.Image(source="upload",type="filepath", label="Imagen")],
  outputs= [gr.outputs.Label(num_top_classes=2, label='Retina'), gr.outputs.Label(num_top_classes=4, label='Enfermedad')],
  title = title, description = description, article = article,
  theme = 'peach',
  examples = ['10.png', '82.png', '15.png', '25.png', '48.png', '61.png', '37.png', '631.png', '23.png', '8.png']
)

interface.launch()