#import required libraries-----------------------------------------------
import streamlit as st
from tensorflow import keras
import tensorflow as tf
from PIL import Image
import numpy as np
import cv2
from imutils import perspective
from scipy.spatial import distance as dist



#loading model-----------------------------------------------------------
model=keras.models.load_model("3rdm_att_UNet_50epochs_acc.h5")

#User Interface----------------------------------------------------------
st.header("Segmentation of the Lower Left Third Molar in Panoramic X-ray Images Using Attention U-Net")

examples=["2.png","20.png","31.png"]

def load_image(image_file):
      img = Image.open(image_file)
      return img
        
st.subheader("Upload Dental Panoramic X-ray Image")
image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])

col1, col2, col3 = st.columns(3)
with col1:
    ex=load_image(examples[0])
    st.image(ex,width=200)
    if st.button('Sample 1'):
        image_file=examples[0]

with col2:
    ex1=load_image(examples[1])
    st.image(ex1,width=200)
    if st.button('Sample 2'):
        image_file=examples[1]

with col3:
    ex2=load_image(examples[2])
    st.image(ex2,width=200)
    if st.button('Sample 3'):
        image_file=examples[2]
    
#main--------------------------------------------------------------------
def midpoint(ptA, ptB):
    return ((ptA[0] + ptB[0]) /2 , (ptA[1] + ptB[1]) /2)

def draw_dimensions(orig_image,predict_image,erode_iteration,open_iteration):
    kernel1 =( np.ones((5,5), dtype=np.float32))#
    kernel_sharpening = np.array([[-1,-1,-1], 
                                  [-1,9,-1], 
                                 [-1,-1,-1]])#
                                 
    image = predict_image
    image2 = orig_image 

    image = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel1,iterations=open_iteration )
    image = cv2.filter2D(image, -1, kernel_sharpening)
    image = cv2.erode(image,kernel1,iterations =erode_iteration)
    #image = image.astype("uint8")#

    #image=cv2.cvtColor(image, cv2.COLOR_BAYER_GR2GRAY)#new
    image=cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)#original

    thresh = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]
    labels=cv2.connectedComponents(thresh,connectivity=8)[1]       
    a=np.unique(labels)
    count2=0
    for label in a:
        if label == 0:
            continue
    
        # Create a mask
        mask = np.zeros(thresh.shape, dtype="uint8")
        mask[labels == label] = 255
        # Find contours and determine contour area
        cnts,hieararch = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cnts = cnts[0]
        c_area = cv2.contourArea(cnts)
        
        (x,y),radius = cv2.minEnclosingCircle(cnts)
        rect = cv2.minAreaRect(cnts)
        box = cv2.boxPoints(rect)
        box = np.array(box, dtype="int")    
        box = perspective.order_points(box)
        color1 = (list(np.random.choice(range(150), size=3)))  
        color =[int(color1[0]), int(color1[1]), int(color1[2])]  
        cv2.drawContours(image2,[box.astype("int")],0,color,2) 
        (tl,tr,br,bl)=box
        
        (tltrX,tltrY)=midpoint(tl,tr)
        (blbrX,blbrY)=midpoint(bl,br)
    	# compute the midpoint between the top-left and top-right points,
    	# followed by the midpoint between the top-righ and bottom-right
        (tlblX,tlblY)=midpoint(tl,bl)
        (trbrX,trbrY)=midpoint(tr,br)
    	# draw the midpoints on the image
        cv2.circle(image2, (int(tltrX), int(tltrY)), 5, (255, 0, 0), -1)
        cv2.circle(image2, (int(blbrX), int(blbrY)), 5, (255, 0, 0), -1)
        cv2.circle(image2, (int(tlblX), int(tlblY)), 5, (255, 0, 0), -1)
        cv2.circle(image2, (int(trbrX), int(trbrY)), 5, (255, 0, 0), -1)
        cv2.line(image2, (int(tltrX), int(tltrY)), (int(blbrX), int(blbrY)),color, 2)
        cv2.line(image2, (int(tlblX), int(tlblY)), (int(trbrX), int(trbrY)),color, 2)
        dA = dist.euclidean((tltrX, tltrY), (blbrX, blbrY))
        dB = dist.euclidean((tlblX, tlblY), (trbrX, trbrY))

        global dimA
        dimA = dA*0.08
        global dimB
        dimB = dB*0.08
        cv2.putText(image2, "{:.1f} millimeter".format(dimA),(int(tltrX - 10), int(tltrY - 10)), cv2.FONT_HERSHEY_SIMPLEX,0.65, (0,0,0), 2)
        cv2.putText(image2, "{:.1f} millimeter".format(dimB),(int(trbrX + 10), int(trbrY+10)), cv2.FONT_HERSHEY_SIMPLEX,0.65,(0,0,0), 2)

    return image2  
    
if image_file is not None:

      img=load_image(image_file)
      
      #orig_to_pass=img
      
      st.text("Making A Prediction ....")
      st.image(img,width=850)
      
      
      #orig_to_pass= np.asarray(orig_to_pass)

      img = np.asarray(img) 
      img_prd= cv2.resize(img , (512, 256))
      img_prd = cv2.cvtColor(img_prd, cv2.COLOR_BAYER_GR2GRAY)
      img_prd = np.expand_dims(img_prd, axis=0)
      prediction = model.predict(img_prd)
      prediction=prediction*255
      prediction = prediction.astype("uint8")
      prediction_img=prediction.reshape(256,512)
      #img3 = prediction_img.astype("uint8")
      img2 = np.zeros( ( np.array(prediction_img).shape[0], np.array(prediction_img).shape[1], 3 ) )#
      img2[:,:,0] = prediction_img # same value in each channel
      img2[:,:,1] = prediction_img
      img2[:,:,2] = prediction_img
      img2 = img2.astype("uint8")
      predicted = cv2.resize(img2 , (img.shape[1],img.shape[0]), interpolation=cv2.INTER_LANCZOS4)
      output=draw_dimensions(img,predicted,3,2)

      if output is not None :      
          st.subheader("Predicted Image")  
          st.write(output.shape)
          st.image(output,width=850)

      st.text("DONE")