import os
import gradio as gr
import tensorflow as tf
import numpy as np
import requests
import torch
import heapq
from PIL import Image

from huggingface_hub import snapshot_download
from huggingface_hub import hf_hub_download

show_all = False
def readLines(filename):
    with open(filename,'r') as f:
        return(f.read().splitlines())


labels = readLines('labels.txt')
labels = [x for x in labels if x] #remove blanks
label_count = len(labels)
valid =  readLines('valid.txt')

model_path = snapshot_download(repo_id="eshieh2/jaguarid_pantanal")
model = tf.saved_model.load(f"{model_path}/saved_model")
serving = model.signatures['serving_default']

detector_path = hf_hub_download(repo_id= "eshieh2/jaguarhead",
                                             filename = "jaguarheadv5.pt")
detector = torch.hub.load('ultralytics/yolov5', 'custom', path = detector_path)

topk = 3

def classify_image(in_image):
  if in_image is None:
    return None
  width,height = in_image.size
  heads = detector(in_image)
  masks = [] # tuple of box coords and string
  for head in heads.xyxy[0]:
    x,y,x2,y2,pct,cls = head.numpy()
    w = x2 - x
    h = y2 - y
    inp = in_image.crop((x,y,x2,y2))
    inp = inp.resize((480,480),Image.BILINEAR)
    inp = np.array(inp)
    inp = np.reshape(inp,(-1, 480, 480, 3)).astype(np.float32)
    inp = np.divide(inp,255.0)
    prediction = serving(tf.convert_to_tensor(inp))['output_0']
    prediction = tf.squeeze(prediction)
    pred = {labels[i]: float(prediction[i]) for i in range(label_count)}
    #print(pred)
    rect = (int(x),int(y),int(x2),int(y2))
    if topk is not None:
      top = heapq.nlargest(topk,pred,key=pred.get)
      label = ''
      for t in top:
        if show_all or t.lower() in valid:
          if len(label) != 0:
            label += ", "
          label += f"{t}:{pred[t]:.3f}"
      if len(label)==0:
        label = 'unknown'
      masks.append((rect,label))
    else:
      max_key = max(pred, key=pred.get)
      if show_all or max_key.lower() in valid:
            masks.append((rect,f"{max_key}:{pred[max_key]}"))
      else:
            masks.append((rect,f"unknown",))
  return (in_image,masks)

image = gr.Image(type='pil')
output = gr.AnnotatedImage()
title = "JaguarID AI identification App"
desc = "Identifies the following: Medrosa, Guaraci, Marcela, Bagua, Manath, Margo, Ti, Ousado, Saseka, Patricia, Kasimir. Confidence displayed after the name is on a 0.0 - 1.0 scale. Must have a clear front/side head view."

gr.Interface(fn=classify_image, inputs=image, outputs=output, examples = [["medrosa.jpg"],["guaraci.jpg"],["marcela.jpg"]], title = title, description = desc).launch()