import gradio as gr
from PIL import Image
import numpy as np
from transformers import AutoImageProcessor, AutoModelForImageClassification

processor_pipe = AutoImageProcessor.from_pretrained("jtas/fish_classification")
model_pipe = AutoModelForImageClassification.from_pretrained("jtas/fish_classification")

def classify_pipe(image):
    inputs = processor_pipe(images=image, return_tensors="pt")

    outputs = model_pipe(**inputs)
    logits = outputs.logits

    # Get the predicted label
    predicted_class_idx = logits.argmax(-1).item()
    labels = model_pipe.config.id2label
    pipe_score = np.argmax(logits, dim=1).max().item()
    predicted_label = labels[predicted_class_idx]

    # Return the predicted label and score
    return {"label": predicted_label, "score": pipe_score}

# Create a Gradio interface
iface = gr.Interface(fn=classify_pipe, inputs=gr.Image(), outputs="json", title="Fish Classification")
iface.launch()