import gradio as gr
import tensorflow as tf
import numpy as np
import requests
import torch

from huggingface_hub import snapshot_download
from huggingface_hub import hf_hub_download

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

labels = ["missing"] * label_count
labels = readLines('labels.txt')
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)

def classify_image(in_image):
  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))
    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))['model']
    prediction = tf.squeeze(prediction)
    pred = {labels[i]: float(prediction[i]) for i in range(label_count)}
    #print(pred)
    max_key = max(pred, key=pred.get)
    rect = (int(x),int(y),int(x2),int(y2))
    if max_key.lower() in valid:
      masks.append((rect,f"{max_key}:{pct}"))
    else:
      masks.append((rect,f"unknown",))
  return (in_image,masks)

image = gr.Image(type='pil')
output = gr.AnnotatedImage()

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