import sys
sys.path.append('src/') # Add src to path
from elm import ELMClassifier

import gradio as gr
import joblib
from utils import visual_words

class_names = ['Arnold Schwarzenegger', 'Silvester Stallone', 'Taylor Swift']

examples = ['imgs/arnold01.jpg',
            'imgs/arnold02.jpg',
            'imgs/arnold03.jpg',
            'imgs/stallone01.jpg',
            'imgs/stallone02.jpg',
            'imgs/stallone03.jpg',
            'imgs/taylor01.jpg',
            'imgs/taylor02.jpg',
            'imgs/taylor03.jpg']

bovw = joblib.load('bovw.joblib') # Load the bag of visual words
clf = joblib.load('elm.joblib') # Load the ELM classifier

def predict_image(img):
    
    h, w = img.shape
    img = img.reshape(-1, h, w) # Reshape image to (1, h, w)

    X = visual_words(img, bovw) # Extract visual words

    prediction = clf.predict_proba(X)[0]
    dict_result = {class_names[i]: float(prediction[i]) for i in range(len(class_names))}

    return dict_result

image = gr.components.Image(image_mode='L') # Grayscale
label = gr.components.Label(num_top_classes=len(class_names))

demo = gr.Interface(fn=predict_image, 
                    inputs=image,
                    outputs=label,
                    examples=examples, 
                    title="Bag of Visual Words Classifier",
                    description="Image classification using Bag of Visual Words (BoVW) and Extreme Learning Machine (ELM) classifier.",)


demo.launch()