Update app.py

app.py

@@ -6,18 +6,40 @@ from vit_model_test import CustomModel
 model = CustomModel()
 
 def predict(image: Image.Image):
-    # Get predictions from the model
-    label, confidence = model.predict(image)
-    
-    # Determine the result based on the label
-    if label == 1:
-        result = "AI image"
-    else:
-        result = "Real image"
-        
-    return result, f"Confidence: {confidence:.2f}%"
+    try:
+        label, confidence = model.predict(image)
+        result = "AI image" if label == 1 else "Real image"
+        return result, f"Confidence: {confidence:.2f}%"
+    except Exception as e:
+        return "Error in prediction", str(e)
 
-# Define the Gradio interface with updated API
+# Custom HTML and CSS to replace the logo with a video
+custom_html = """
+<style>
+    /* Hide the default logo */
+    .gradio-logo {
+        display: none;
+    }
+    /* Center the video container */
+    .loading-container {
+        text-align: center;
+        margin-top: 20px;
+    }
+    .loading-container video {
+        width: 320px; /* Adjust video size */
+        height: auto; /* Maintain aspect ratio */
+    }
+</style>
+<div class="loading-container">
+    <video autoplay muted loop>
+        <source src="load_screen.mp4" type="video/mp4">
+        Your browser does not support the video tag.
+    </video>
+    <div>Processing, please wait...</div>
+</div>
+"""
+
+# Define the Gradio interface
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
@@ -26,6 +48,8 @@ demo = gr.Interface(
     description="Upload an image to classify it using the Vision Transformer model."
 )
 
-# Launch the Gradio interface
+# Inject the custom HTML to show the video instead of the logo
+demo.load(custom_html)
 
+# Launch the Gradio interface
 demo.launch()