update gradio interface and inference logic

app.py

@@ -1,19 +1,46 @@
 import gradio as gr
 from PIL import Image
 import numpy as np
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline
 
-processor = AutoImageProcessor.from_pretrained("jtas/fish_classification")
-classifier = AutoModelForImageClassification.from_pretrained("jtas/fish_classification")
+classifier = pipeline(model="jtas/fish_classification")
+model = AutoModelForImageClassification.from_pretrained("jtas/fish_classification")
+
+fish_classes = model.config.id2label
 
 def fish_classification(image):
-    inputs = processor(images=image, return_tensors="pt")
+    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"],
+]
 
-    outputs = classifier(**inputs)
+supported_classes_html = ""
+for class_id, class_name in fish_classes.items():
+    supported_classes_html += f"<p><b>{class_id}</b>: {class_name}</p>"
 
-    return outputs
+label = gr.components.Label()
 
-label_fish = gr.components.Json()
+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 = gr.Interface(fn=fish_classification, inputs=gr.Image(), outputs=label_fish, title="Fish Classification")
-iface.launch()
+iface.launch()