Update app.py

app.py

@@ -1,26 +1,43 @@
 import gradio as gr
 from PIL import Image
-from vit_model_test import CustomModel
+from vit_model_test import Custom_VIT_Model  # Ensure you import the correct class
 
 # Initialize the model
-model = CustomModel()
+model = Custom_VIT_Model()
+
+# Variable to store the last prediction result
+last_prediction = None
 
 def predict(image: Image.Image):
-        label, confidence = model.predict(image)
-        result = "AI image" if label == 1 else "Real image"
-        return result, f"Confidence: {confidence:.2f}%"
+    global last_prediction
+    label, confidence = model.predict(image)
+    result = "AI image" if label == 1 else "Real image"
+    last_prediction = (image, label)  # Store the image and prediction label
+    return result, f"Confidence: {confidence:.2f}%"
 
+def report_feedback():
+    if last_prediction is not None:
+        image, predicted_label = last_prediction
+        correct_label = 1 if predicted_label == 0 else 0  # Invert the label
+        model.add_data(image, correct_label)  # Add incorrect prediction to model
+        return "Feedback recorded. Thank you!"
+    return "No prediction available to report."
 
-    # Define the Gradio interface
+# Define the Gradio interface for prediction
 demo = gr.Interface(
-fn=predict,
-inputs=gr.Image(type="pil"),
-outputs=[gr.Textbox(), gr.Textbox()],
-title="Vision Transformer Model",
-description="Upload an image to classify it using the Vision Transformer model.",
-theme = gr.themes.Soft()
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Textbox(label="Prediction"), gr.Textbox(label="Confidence")],
+    title="Vision Transformer Model",
+    description="Upload an image to classify it using the Vision Transformer model.",
+    theme=gr.themes.Soft()
 )
 
+# Define the feedback button
+feedback_button = gr.Button("The model was wrong")
+feedback_button.click(report_feedback)
+
 # Launch the Gradio interface
 if __name__ == "__main__":
     demo.launch(share=True)
+    feedback_button.launch()