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')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, dataset_name="ofwat_cleaner_loop", private=True, 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}

def flag_feedback(image, option, flag_status):
    # Perform flagging action here using hf_writer
    hf_writer.flag((image, option))

    # Update the flag status to indicate feedback has been submitted
    flag_status.update("Feedback submitted. Thank you!")
    return flag_status

# Create a state variable for the flag status
flag_status = gr.State("")

# Create the Gradio interface
iface = gr.Interface(
    classify_image,
    inputs=[gr.Image(), gr.Radio(["obstruction", "no_obstruction"])],
    outputs=[gr.Label(), gr.Textbox(label="Flag Status", value=flag_status)],
    examples=[
        ["examples/CS.jpg"],
        ["examples/GI.jpg"],
        ["examples/PP.jpg"]
    ],
    description="Upload an image to view a classification demonstration...",
    title="Sewer Obstruction Classification with AI by Subterra",
    allow_flagging="manual",
    flagging_options=["obstruction", "no_obstruction"],
    flagging_callback=lambda image, option: flag_feedback(image, option, flag_status)
)

# Launch the interface
iface.launch()