from datasets import load_dataset
import gradio as gr
from transformers import pipeline
from PIL import Image

# Load the dataset
dataset = load_dataset("neuraxcompany/coin_classification")

# Initialize the image classification pipeline with the ViT model
classifier = pipeline("image-classification", model="google/vit-base-patch16-224")

# Coin identification function using dataset labels
def identify_coin(image):
    # Convert image to PIL format if needed
    if isinstance(image, Image.Image):
        img = image
    else:
        img = Image.fromarray(image)

    # Run classification
    results = classifier(img)

    # Format and match dataset labels
    coin_labels = {entry["CoinType"]: entry["Side"] for entry in dataset["train"]}
    predictions = {res["label"]: round(res["score"], 4) for res in results}

    # Match predictions with dataset labels
    matched_labels = {coin_labels.get(label, label): score for label, score in predictions.items()}

    return matched_labels

# Greeting function
def greet(name):
    return f"Hello {name}!!"

# Create the Gradio interface with separate tabs
greet_interface = gr.Interface(fn=greet, inputs="text", outputs="text", title="Greeting")
coin_identifier_interface = gr.Interface(fn=identify_coin, inputs=gr.Image(), outputs=gr.Label(), title="Coin Classifier")

demo = gr.TabbedInterface([greet_interface, coin_identifier_interface])

# Launch the app
demo.launch()