Create demo

app.py

@@ -0,0 +1,43 @@
+import gradio as gr
+from utils import model_initialization, prediction
+from PIL import Image
+from typing import Dict, Any
+
+
+def gradio_interface(image: Image.Image) -> Dict[str, Any]:
+    """
+    Perform image classification using a pre-trained model.
+
+    Args:
+        image (Image.Image): The input image uploaded by the user.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the classification result with the
+        most promising label and confidence score.
+    """
+    # Initialize the pre-trained pipeline
+    pipe = model_initialization()
+    
+    # Perform prediction on the uploaded image
+    result = prediction(pipe, image)
+
+    return result
+
+
+# Define the Gradio interface
+demo = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Image(type="pil", label="Upload Image"),  # Accepts PIL Image input
+    outputs=gr.JSON(label="Prediction Details"),  # Outputs as JSON
+    title="RESNET WILL NEVER DIE. Image Classification with ResNet-18",
+    description=(
+        "Welcome to the Image Classification Demo! Upload an image to classify it using"
+        "ResNet-18 model. The model will predict the most likely label along with its confidence score."
+    ),
+    theme="soft",
+    examples=[["artifacts/ball.png"], ["artifacts/panda.jpg"]],
+)
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+transformers==4.46.3
+torch==2.5.1

utils.py

@@ -0,0 +1,47 @@
+from transformers import pipeline, Pipeline
+from functools import lru_cache
+from typing import Optional, Dict, Any
+import numpy as np
+
+
+@lru_cache
+def model_initialization(task: str = "image-classification", model_name: str = "microsoft/resnet-18") -> Pipeline:
+    """
+    Initialize the Hugging Face pipeline for a specified task and model.
+
+    Args:
+        task (str): The task type, e.g., "image-classification".
+        model_name (str): The name or path of the model to use.
+
+    Returns:
+        Pipeline: A Hugging Face pipeline object ready for inference.
+    """
+    pipe = pipeline(task, model=model_name)
+
+    return pipe
+
+
+def prediction(pipe: Pipeline, img: np.ndarray) -> Optional[Dict[str, Any]]:
+    """
+    Perform image classification on the given image using the specified pipeline.
+
+    Args:
+        pipe (Pipeline): The initialized hf pipeline object.
+        img (np.ndarray): The image to classify.
+
+    Returns:
+        Optional[Dict[str, Any]]: A dictionary containing the most promising label and its confidence score,
+        or None if no results are returned.
+    """
+    results = pipe(img)
+    results.sort(key=lambda x: x["score"], reverse=True)
+
+    if not results:
+        return None
+    
+    response = {
+        "most_promising_label": results[0]["label"],
+        "confidence": round(results[0]["score"], 2)
+    }
+
+    return response