app.py

import gradio as gr
from transformers import pipeline
from PIL import Image
import os

# Initialize the pipeline with your model
pipe = pipeline("image-classification", model="SubterraAI/ofwat_cleaner_classification")
HF_TOKEN = os.environ.get('HF_TOKEN')
print(HF_TOKEN)
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "SubterraAI/ofwat_cleaner_loop", separate_dirs=True)

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"]
    ],
    description="Upload an image to view a classification demonstration leveraging the dataset/library of images collected by WRc & Unitied Utitlies during The Water Services Regulation Authority (OFWAT) Innovation Challenge – Artificial Intelligence and Sewers. Not only can you see the initial classification, but you as the user can also inform us if the classification is correct. Your response will be used to retrain this model. The team at Subterra would like to thank all of those involved in collecting this dataset as we hope that other groups will use it to further advance technology solutions for the water industry.",
    title="Sewer Obstruction Classification with AI by Subterra",
    allow_flagging="manual",
    flagging_options=["obstruction", "no_obstruction"],
    flagging_callback=hf_writer

)

# Launch the interface
iface.launch()