from huggingface_hub import from_pretrained_keras
import gradio as gr
from PIL import Image
import tensorflow as tf
import numpy as np
import os

# load model
model = tf.keras.models.load_model('./tf_model.h5') # from keras-io/semantic-segmentation

# gradio components
inputs = gr.inputs.Image()
outputs = gr.outputs.Image()

# inference fn
def predict(inputs):
  # convert img to numpy array, resize and normalize to make the prediction
  img = np.array(inputs)

  im = tf.image.resize(img, (128, 128))
  im = tf.cast(im, tf.float32) / 255.0
  pred_mask = model.predict(im[tf.newaxis, ...])
  
  # take the best performing class for each pixel
  # the output of argmax looks like this [[1, 2, 0], ...]
  pred_mask_arg = tf.argmax(pred_mask, axis=-1)

  labels = []
  
  # convert the prediction mask into binary masks for each class
  binary_masks = {}
  mask_codes = {}
  
  # when we take tf.argmax() over pred_mask, it becomes a tensor object
  # the shape becomes TensorShape object, looking like this TensorShape([128]) 
  # we need to take get shape, convert to list and take the best one
  
  rows = pred_mask_arg[0][1].get_shape().as_list()[0]
  cols = pred_mask_arg[0][2].get_shape().as_list()[0]
  
  for cls in range(pred_mask.shape[-1]):

      binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class
      
      for row in range(rows):

          for col in range(cols):

              if pred_mask_arg[0][row][col] == cls:
                  
                  binary_masks[f"mask_{cls}"][row][col] = 1
              else:
                  binary_masks[f"mask_{cls}"][row][col] = 0

      mask = binary_masks[f"mask_{cls}"]
      mask *= 255
      img = Image.fromarray(mask.astype(np.int8), mode="L")
      return img

  
gr.Interface(predict, inputs = inputs, outputs = outputs).launch()