import gradio as gr
import torch
import time
from fastai.vision.all import load_learner
from PIL import Image

# Load the exported model
model = load_learner("model.pkl")

# Function to classify an image
def classify_images(imgs):
    start_time = time.time()
    results = []
    for img in imgs:
        # Convert gradio image to PIL Image
        pil_img = Image.fromarray(img.astype('uint8'), 'RGB')

        # Perform inference
        pred_class, pred_idx, pred_probs = model.predict(pil_img)

        # Format output
        output = f"Image Name: {img.name} - Category: {pred_class}"
        results.append(output)

    # Calculate total inference time
    inference_time = time.time() - start_time

    # Append total inference time to results
    results.append(f"Total Inference Time: {inference_time:.2f} seconds")

    return results

# Create Gradio interface
input_component = gr.inputs.Image(label="Upload Image", type="file", multiple_files=True)
output_component = gr.outputs.Textbox(label="Classification Results")
interface = gr.Interface(fn=classify_images, inputs=input_component, outputs=output_component, title="Image Classifier")

# Launch the Gradio interface
interface.launch()