import streamlit as st  #Web App
from PIL import Image, ImageOps #Image Processing
import time
from unittest import result
from icevision.all import *
from icevision.models.checkpoint import *


st.sidebar.image("./logo.png")
st.sidebar.header("ATK-OCR classification (AOC) Webapp.")
def load_image(image_file):
    img = Image.open(image_file)
    return img



activities = ["Detection (วิเคราะห์โรค)", "About (เกี่ยวกับ)"]
choice = st.sidebar.selectbox("Select option.. (เลือกโหมด)",activities)

#set default size as image_scale x image_scale
def img_resize(input_path,img_size): # padding
  desired_size = img_size
  im = Image.open(input_path)
  im = ImageOps.exif_transpose(im) # fix image rotating
  width, height = im.size # get img_input size
  if (width == image_scale) and (height == image_scale):
    new_im = im
  else:
    #im = im.convert('L') #Convert to gray
    old_size = im.size  # old_size[0] is in (width, height) format
    ratio = float(desired_size)/max(old_size)
    new_size = tuple([int(x*ratio) for x in old_size])
    im = im.resize(new_size, Image.ANTIALIAS)
    new_im = Image.new("RGB", (desired_size, desired_size))
    new_im.paste(im, ((desired_size-new_size[0])//2,
                        (desired_size-new_size[1])//2))

  return new_im

checkpoint_path = "./AOC_weight_97.4.pth"

checkpoint_and_model = model_from_checkpoint(checkpoint_path, 
    model_name='ross.efficientdet', 
    backbone_name='tf_d2',
    img_size=384, 
    is_coco=False)

model_type = checkpoint_and_model["model_type"]
backbone = checkpoint_and_model["backbone"]
class_map = checkpoint_and_model["class_map"]
img_size = checkpoint_and_model["img_size"]
#model_type, backbone, class_map, img_size

model = checkpoint_and_model["model"]

device=next(model.parameters()).device

img_size = checkpoint_and_model["img_size"]
valid_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(img_size), tfms.A.Normalize()])

image_scale = 768 # change image input / ouput size here

def get_detection(img_path):
 
  #Get_Idcard_detail(file_path=img_path)
  img = Image.open(img_path)
  img = ImageOps.exif_transpose(img) # fix image rotating
  width, height = img.size # get img_input size
  if (width == image_scale) and (height == image_scale):
    pred_dict  = model_type.end2end_detect(img, valid_tfms, model, class_map=class_map, detection_threshold=0.6)
  else:
    #im = im.convert('L') #Convert to gray
    old_size = img.size  # old_size[0] is in (width, height) format
    ratio = float(image_scale)/max(old_size)
    new_size = tuple([int(x*ratio) for x in old_size])
    img = img.resize(new_size, Image.ANTIALIAS)
    new_im = Image.new("RGB", (image_scale, image_scale))
    new_im.paste(img, ((image_scale-new_size[0])//2,
                        (image_scale-new_size[1])//2))
    pred_dict  = model_type.end2end_detect(new_im, valid_tfms, model, class_map=class_map, detection_threshold=0.6)

    if (len(pred_dict['detection']['labels'])) > 0:
      num_count = 1
      for i,j in zip(pred_dict['detection']['labels'], pred_dict['detection']['scores']):
        labels = i
        acc = round((j*100),2)
        if labels == "Neg":
          labels = "Negative"
          st.success(f"{num_count}. '{labels}' with {acc} % confidence.")
        elif labels == "Pos":
          labels = "Positive"
          st.error(f"{num_count}. '{labels}' with {acc} % confidence.")
        num_count += 1
    else: # when not found imgs
      st.warning("Not found 'Antigen Test Kit' image ; please recheck and try again !!")

def get_img_detection(img_path):
   
  #Get_Idcard_detail(file_path=img_path)
  img = Image.open(img_path)
  img = ImageOps.exif_transpose(img) # fix image rotating
  width, height = img.size # get img_input size
  if (width == image_scale) and (height == image_scale):
    new_im = img
  else:
    #im = im.convert('L') #Convert to gray
    old_size = img.size  # old_size[0] is in (width, height) format
    ratio = float(image_scale)/max(old_size)
    new_size = tuple([int(x*ratio) for x in old_size])
    img = img.resize(new_size, Image.ANTIALIAS)
    new_im = Image.new("RGB", (image_scale, image_scale))
    new_im.paste(img, ((image_scale-new_size[0])//2,
                        (image_scale-new_size[1])//2))
  
  pred_dict  = model_type.end2end_detect(new_im, valid_tfms, model, class_map=class_map, detection_threshold=0.6)


  return pred_dict['img']






if choice =='About (เกี่ยวกับ)' :
    st.header("About... (เกี่ยวกับ)")

    st.subheader("AOC คืออะไร ?")
    st.write("- เป็นระบบที่สามารถคัดกรองผลตรวจเชื้อของ COVID-19 ได้ผ่าน ที่ตรวจ ATK (Antigen Test Kit) ควบคู่กับบัตรประชาชน จากรูปภาพได้โดยอัตโนมัติ")

    st.subheader("AOC ทำอะไรได้บ้าง ?")
    st.write("- ตรวจจับผลตรวจ ATK (Obj detection)")
    st.write("- ตรวจจับชื่อ-นามสกุล (OCR)")
    st.write("- ตรวจจับเลขบัตรประชาชน (OCR)")

    st.subheader("AOC ดีกว่ายังไง ?")
    st.write("จากผลที่ได้จากการเปรียบเทียบกันระหว่าง model (AOC) กับ คน (Baseline) จำนวน 30 ภาพ / คน ได้ผลดังนี้")
    st.image("./acc_table.png")
    st.write("จากผลที่ได้สรุปได้ว่า ส่วนที่ผ่าน Baseline และมีประสิทธิภาพดีกว่าคัดกรองด้วยคนคือ ผลตรวจ ATK ได้ผลที่ 100 %, เลขบัตรประชน ได้ผลที่ 100 % และ ความเร็วในการคัดกรอง ได้ผลที่ 4.84 วินาที ซึ่งมีความเร็วมากกว่า 81% เมื่อเทียบกับคัดกรองด้วยคน ถือว่ามีประสิทธิภาพที่สูงมากในการคัดกรอง และ มีประสิทธิภาพมากกว่าการคัดแยกด้วยมนุษย์")
    st.write("** ความเร็วที่โมเดลทำได้อาจไม่ตรงตามที่ deploy บนเว็บ เนื่องจากในเว็บ ไม่มี GPU ในการประมวลผลอาจทำให้โมเดลใช้เวลาในการประมวลที่นานกว่าตอนใช้ GPU")


    st.subheader("คำแนะนำในการใช้งาน")
    st.write("- ในการใช้งานให้ถ่ายรูปภาพบัตรประชาชนในแนวตั้งเท่านั้น เนื่องจากถ้าเป็นแนวอื่นอาจทำให้การตรวจจับคลาดเคลื่อนเอาได้")#3
    st.write("- ภาพไม่ควรมีแสงที่สว่างมากเกืนไป และ มืดเกินไป มิฉะนั้นอาจทำให้การตรวจจับคลาดเคลื่อนเอาได้")#4
    st.write("- ภาพไม่ควรที่จะอยู่ไกลเกินไป และ ควรมีความชัด มิฉะนั้นอาจทำให้การตรวจจับคลาดเคลื่อน หรือ ไม่สามารถตรวจจับได้")#5

    st.subheader("รายละเอียดเพิ่มเติม")
    st.write('[Medium blog](https://medium.com/@mjsalyjoh/atk-ocr-classification-aoc-%E0%B8%A3%E0%B8%B0%E0%B8%9A%E0%B8%9A%E0%B8%84%E0%B8%B1%E0%B8%94%E0%B8%81%E0%B8%A3%E0%B8%AD%E0%B8%87%E0%B8%9C%E0%B8%A5%E0%B8%95%E0%B8%A3%E0%B8%A7%E0%B8%88-atk-%E0%B9%81%E0%B8%A5%E0%B8%B0-%E0%B8%9A%E0%B8%B1%E0%B8%95%E0%B8%A3%E0%B8%9B%E0%B8%A3%E0%B8%B0%E0%B8%8A%E0%B8%B2%E0%B8%8A%E0%B8%99-fa32a8d47599)')
    st.write('[Github Link](https://github.com/Tanaanan/AOC_ATK_OCR_Classification)')

    
     
       
elif choice == "Detection (วิเคราะห์โรค)":
    st.header(" Antigen test kit + Identification card detector.")

    image = st.file_uploader(label = "upload ATK + Idcard img here.. OwO",type=['png','jpg','jpeg'])
    if image is not None:
        new_img = img_resize(image, 1280)
        st.image(get_img_detection(image))
        with st.spinner("🤖 On Working... "):
            t1 = time.perf_counter()
            st.subheader("Detection result...")
            get_detection(image)
            t2 = time.perf_counter()
            st.write('time taken to run: {:.2f} sec'.format(t2-t1))
    else:
        st.write("## Waiting for image..")
        st.image('atk_idcard.jpeg')

    st.warning("""[On debugging..] ในตอนนี้ระบบ AOC สามารถใช้ได้ในการตรวจจับ ATK เพียงอย่างเดียว (กำลังอยู่ในขั้นตอนการพัฒนาให้สามารถใช้งานควบคู่กับระบบ OCR บัตรประชาชน)""")

    st.caption("Made by Tanaanan .M")


st.sidebar.subheader('More image for test..')
st.sidebar.write('[Github img test set.](https://github.com/Tanaanan/AOC_ATK_OCR_Classification/tree/main/test_set(img))')
    









st.sidebar.markdown('---')
st.sidebar.subheader('Recomend / Issues report..')
st.sidebar.write('[Google form](https://forms.gle/zYpYFKcTpBoFGxN58)')


st.sidebar.markdown('---')
st.sidebar.subheader('Made by Tanaanan .M')
st.sidebar.write("Contact : mjsalyjoh@gmail.com")