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

classifier = pipeline(model="jtas/fish_classification")
model = AutoModelForImageClassification.from_pretrained("jtas/fish_classification")

fish_classes = model.config.id2label

def fish_classification(image):
    pil_img = Image.fromarray(np.uint8(image))

    fish_prediction = classifier(pil_img)
    class_probs = {str(pred["label"]): pred["score"] for pred in fish_prediction}

    return class_probs

sample_images = [
    ["img/Rastrelliger Faughni.jpg", "Rastrelliger Faughni"],
    ["img/Chanos Chanos.jpg", "Chanos Chanos"],
    ["img/Eleutheronema Tetradactylum.jpeg", "Eleutheronema Tetradactylum"],
    ["img/Johnius Trachycephalus.jpg", "Johnius Trachycephalus"],
    ["img/Nibea Albiflora.jpeg", "Nibea Albiflora"],
    ["img/Oreochromis Mossambicus.jpg", "Oreochromis Mossambicus"],
    ["img/Oreochromis Niloticus.png", "Oreochromis Niloticus"],
    ["img/Upeneus Moluccensis.jpg", "Upeneus Moluccensis"],
]

supported_classes_html = ""
for class_id, class_name in fish_classes.items():
    supported_classes_html += f"<p><b>{class_id}</b>: {class_name}</p>"

label = gr.components.Label()

iface = gr.Interface(
    fn=fish_classification,
    inputs=gr.Image(label="Upload an image"),
    examples=sample_images,
    outputs=label,
    title="Fish Classification",
    description="This web app classifies fish in an image. Upload an image of a fish to see the predicted class probabilities.",
    article="<div style='margin-top: 20px;'><h2>Supported Fish Classes</h2>" + supported_classes_html + "</div>",
)

iface.launch()