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

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

# Function to classify an image
def classify_images(imgs):
    [print(x) for x in imgs]
    start_time = time.time()
    results = []
    for img in imgs:
        # Convert gradio image to PIL Image
        #img = PILImage.create(img)

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

        # Format output
        output = f"Image Name: {Path(img).stem} - 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.File(label="Upload Image", file_count='multiple')
output_component = gr.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()