from turtle import title
import gradio as gr
from transformers import pipeline
from PIL import Image

# Initialize the pipeline with your model
pipe = pipeline("image-classification", model="SubterraAI/ofwat_material_classification")

def classify_image(image):
    # Convert the input image to PIL format
    PIL_image = Image.fromarray(image).convert('RGB')
    
    # Classify the image using the pipeline
    res = pipe(PIL_image)
    
    # Extract labels and scores
    return {dic["label"]: dic["score"] for dic in res}


# Create the Gradio interface
iface = gr.Interface(
    classify_image, 
    "image", 
    "label", 
    examples=[
        ["examples/CS.jpg"], 
        ["examples/GI.jpg"],
        ["examples/PP.jpg"],
        ["examples/RC.jpg"]
    ],
    description="Upload an image to classify its material.",
    title="Material Classification with AI by Subterra"
)

# Launch the interface
iface.launch()