import gradio as gr
import os
from transformers import pipeline
from pathlib import Path
from PIL import Image
import numpy as np

example_imgs = ["examples/img0.jpg", 
                "examples/img1.jpg",
                "examples/img2.jpg",
                "examples/img3.jpg"]

pipe = pipeline("image-classification", model="arnaucas/wildfire-classifier")

def inference(image):
    image = Image.fromarray(np.uint8(image)).convert('RGB')
    output = pipe(image)
    result = {item['label']: item['score'] for item in output}
    return result

gr.Interface(
    fn=inference,
    title="Wildfire Detection",
    description = "Predict whether an image contains wildfire or not",
    inputs="image",
    examples=example_imgs,
    outputs=gr.Label(),
    cache_examples=False,
    theme='earneleh/paris',
    article = "Author: <a href=\"https://www.linkedin.com/in/arnau-castellano/\">Arnau Castellano</a>",
).launch(debug=True, enable_queue=True)