import gradio as gr
import json
from transformers import pipeline


def load_label_to_name_mapping(json_file_path):
    """Load the label-to-name mapping from a JSON file."""
    with open(json_file_path, 'r') as f:
        mapping = json.load(f)
    return {int(k): v for k, v in mapping.items()}

def infer_flower_name(classifier, image):
    """Perform inference on an image and return the flower name."""
    # Perform inference
    # Load the model checkpoint for inference
    
    result = classifier(image)
    # Get the label from the inference result
    label = result[0]['label'].split('_')[-1]  # The label is usually in the format 'LABEL_#'
    label = int(label)
    
    # Map the integer label to the flower name
    json_file_path = 'label_to_name.json'
    label_to_name = load_label_to_name_mapping(json_file_path)
    flower_name = label_to_name.get(label, "Unknown")
    
    return flower_name

def predict(flower): # would call a model to make a prediction on an input and return the output.
    classifier = pipeline("image-classification", model="checkpoint-160")
    flower_name = infer_flower_name(classifier, flower)
    return flower_name


description = "Upload an image of a flower and discover its species!"
title = "Bloom Classifier"
examples = ["examples/example.jpg", "examples/image_00293.jpg","examples/image_02828.jpg"]
demo = gr.Interface(fn=predict, 
                    inputs=gr.Image(type="pil"), 
                    outputs=gr.Label(num_top_classes=3),
                    description=description,
                    title = title,
                    live = False,
                    share=True,
                    examples=examples)

demo.launch()