### -------------------------------- ###
###            libraries             ###
### -------------------------------- ###
import gradio as gr
import numpy as np
import os
from tensorflow.keras.models import load_model

### -------------------------------- ###
###           model loading          ###
### -------------------------------- ###
model = load_model('model.h5') # single file model from colab

## --------------------------------- ###
###     reading: categories.txt      ###
### -------------------------------- ###
labels = ['please upload categories.txt' for i in range(10)] # placeholder

if os.path.isfile("categories.txt"):
    # open categories.txt in read mode
    categories = open("categories.txt", "r")
    labels = categories.readline().split()

## --------------------------------- ###
###     rendering: info.html         ###
### -------------------------------- ###
# borrow file reading functionality from reader.py
# info =
description = "A Hugging Space demo created by datasith"
title = "Cast parts: Deffective or Okay?"
# css = \
# '''
#   .div {
#     border: 2px solid black;
#     margin: 10px;
#     padding: 5%;
#   }
#   ul {
#     display: inline-block;
#     text-align: left;
#   }
#   img {
#     display: block;
#     margin: auto;
#   }
#   .description {
#     text-align: center;
#   }
# '''

article = \
'''
Deffective or Okay? Demo app including a binary classification model for casted parts
This is a test project to get familiar with Hugging Face! 
The space includes the necessary files for everything to run smoothly on HF's `Spaces`:
- `app.py`
- `reader.py`
- `requirements.txt`
- `model.h5` (TensorFlow/Keras)
- `categories.txt`
- `info.txt`
The data used to train the model is available as [Kaggle dataset](https://www.kaggle.com/datasets/ravirajsinh45/real-life-industrial-dataset-of-casting-product).
The space was inspired by @Isabel's wonderful [cat or pug](https://huggingface.co/spaces/isabel/pug-or-cat-image-classifier) one. Enjoy!d
'''


### -------------------------------- ###
###        interface creation        ###
### -------------------------------- ###
samples = ['defective.jpeg', 'okay.jpeg']

def preprocess(image):
  image = np.array(image) / 255
  image = np.expand_dims(image, axis=0)
  return image
  
def predict_image(image):
  pred = model.predict(preprocess(image))
  results = {}
  for row in pred:
    for idx, item in enumerate(row):
      results[labels[idx]] = float(item)
  return results

# generate img input and text label output
image = gr.inputs.Image(shape=(300, 300), label="Upload Your Image Here")
label = gr.outputs.Label(num_top_classes=len(labels))

# generate and launch interface
interface = gr.Interface(fn=predict_image, inputs=image, outputs=label, article=article, theme='default', title=title, allow_flagging='never', description=description, examples=samples)
interface.launch()