Improvements in message feedback

app.py

@@ -1,4 +1,3 @@
-from turtle import title
 import gradio as gr
 from transformers import pipeline
 from PIL import Image
@@ -8,7 +7,7 @@ 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)
+hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "ofwat_cleaner_loop", private=True, separate_dirs=True)
 
 
 def classify_image(image):
@@ -40,4 +39,4 @@ iface = gr.Interface(
 )
 
 # Launch the interface
-iface.launch()
+iface.launch()

app_save.py

@@ -0,0 +1,50 @@
+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()