# -*- coding: utf-8 -*-
"""threadcheckerv1_gui.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/145qYaJaBGKmbGoSNYSsFlMASN1PUKEgF
"""

import numpy as np
from PIL import Image
import gradio as gr
import huggingface_hub
from tensorflow.keras.models import load_model

# Load model from Hugging Face Hub
repo_id = "ddecosmo/thread_checker_v1"
model_filename = "thread_checker_model.keras"
model_path = huggingface_hub.hf_hub_download(repo_id=repo_id, filename=model_filename)
model = load_model(model_path)

# Example images (replace with actual files in your repo if desired)
example_images = [
    "0.125_ex.jpg",
    "0.25_ex.jpg",
    "0.375_ex.jpg"
]

def predict_image(image):
    """
    Predicts the class of an image using the loaded Keras model and returns
    confidence scores for all classes and the final determination.
    """
    img_width, img_height = model.input_shape[1:3]
    image = image.resize((img_width, img_height))
    image = np.array(image).astype("float32") / 255.0
    image = np.expand_dims(image, axis=0)

    predictions = model.predict(image)
    confidence_scores = predictions[0]

    predicted_class_index = np.argmax(confidence_scores)
    class_labels = ["0.125", "0.25", "0.375"]
    final_determination = class_labels[predicted_class_index]

    return (
        float(confidence_scores[0]),
        float(confidence_scores[1]),
        float(confidence_scores[2]),
        final_determination,
    )

iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Number(label="Confidence (0.125)"),
        gr.Number(label="Confidence (0.25)"),
        gr.Number(label="Confidence (0.375)"),
        gr.Textbox(label="Final Determination"),
    ],
    examples=example_images,
)

if __name__ == "__main__":
    iface.launch()