Update app.py

app.py

@@ -1,6 +1,26 @@
-import os
+import gradio as gr
+from PIL import Image
+from vit_model_test import CustomModel
 
-if os.path.exists("https://cdn-uploads.huggingface.co/production/uploads/66d6f1b3b50e35e1709bfdf7/x7Ud8PO9QPfmrTvBVcCKE.mp4"):
-    print("The video file exists.")
-else:
-    print("The video file does not exist.")
+# Initialize the model
+model = CustomModel()
+
+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}%"
+    
+
+# Define the Gradio interface
+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."
+)
+
+
+# Launch the Gradio interface
+demo.launch()