Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 from PIL import Image
-from vit_model_test import CustomModel  # Ensure you import the correct class
+from vit_model_test import CustomModel  # Ensure this is the correct import for your model
 from vit_Training import Custom_VIT_Model
 
 # Initialize the model
-model = Custom_VIT_Model()
+model = CustomModel()
 
 # Variable to store the last prediction result
 last_prediction = None
@@ -13,7 +13,7 @@ def predict(image: Image.Image):
     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
+    last_prediction = (image, label)  # Store the image and label for feedback
     return result, f"Confidence: {confidence:.2f}%"
 
 def report_feedback():
@@ -21,29 +21,26 @@ 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
+        model.add_data(image, correct_label)  # Pass the incorrect prediction to the model
         return "Feedback recorded. Thank you!"
     return "No prediction available to report."
 
-# Define the Gradio interface for prediction and feedback
-def main():
-    with gr.Blocks() as demo:
-        gr.Markdown("### Vision Transformer Model")
-        gr.Markdown("Upload an image to classify it using the Vision Transformer model.")
-        
-        image_input = gr.Image(type="pil", label="Upload Image")
-        prediction_output = gr.Textbox(label="Prediction", interactive=False)
-        confidence_output = gr.Textbox(label="Confidence", interactive=False)
-        feedback_output = gr.Textbox(label="Feedback Status", interactive=False)
-        
-        submit_btn = gr.Button("Submit")
-        feedback_btn = gr.Button("The model was wrong")
-        
-        submit_btn.click(predict, inputs=image_input, outputs=[prediction_output, confidence_output])
-        feedback_btn.click(report_feedback, outputs=feedback_output)
-
-    # Launch the Gradio interface
-    demo.launch(share=True)
+# Define the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("### Vision Transformer Model")
+    gr.Markdown("Upload an image to classify it using the Vision Transformer model.")
+    
+    image_input = gr.Image(type="pil", label="Upload Image")
+    prediction_output = gr.Textbox(label="Prediction", interactive=False)
+    confidence_output = gr.Textbox(label="Confidence", interactive=False)
+    feedback_output = gr.Textbox(label="Feedback Status", interactive=False)
 
+    submit_btn = gr.Button("Submit")
+    feedback_btn = gr.Button("The model was wrong")
+
+    submit_btn.click(predict, inputs=image_input, outputs=[prediction_output, confidence_output])
+    feedback_btn.click(report_feedback, outputs=feedback_output)
+
+# Launch the Gradio interface
 if __name__ == "__main__":
-    main()
+    demo.launch(share=True)