Upload main.py

app/main.py

@@ -0,0 +1,69 @@
+import gradio as gr
+import os
+import sys
+
+# Add project root to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from src.inference.predictor import VisionGuardPredictor
+
+# 1. Load Model
+model_path = "models_saved/dinov2_best.pt"
+print(f"⏳ Loading VisionGuard AI ({model_path})...")
+
+try:
+    predictor = VisionGuardPredictor(model_path)
+    print("✅ System Ready.")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    sys.exit(1)
+
+# 2. Logic
+def analyze_image(image):
+    if image is None:
+        return None, None, "Please upload an image."
+    
+    temp_path = "temp_analysis.jpg"
+    image.save(temp_path)
+    
+    try:
+        # Run Prediction
+        result = predictor.predict(temp_path)
+        
+        summary = (
+            f"Verdict: {result['verdict']}\n"
+            f"Confidence: {result['confidence']}%"
+        )
+        
+        # Return: Label Dict, Heatmap Image, Summary Text
+        return result['probabilities'], result['heatmap'], summary
+        
+    except Exception as e:
+        return None, None, f"Error: {str(e)}"
+
+# 3. UI
+with gr.Blocks(title="VisionGuard AI") as demo:
+    gr.Markdown("# 🛡️ VisionGuard AI")
+    gr.Markdown("Upload an image to detect AI artifacts. The **Heatmap** shows which areas triggered the detection.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Source")
+            submit_btn = gr.Button("Analyze Integrity", variant="primary")
+            
+        with gr.Column():
+            # Output 1: Probability
+            label_output = gr.Label(num_top_classes=2, label="Probability")
+            # Output 2: Heatmap
+            heatmap_output = gr.Image(label="Attention Heatmap (X-Ray)")
+            # Output 3: Text
+            info_output = gr.Textbox(label="Verdict")
+            
+    submit_btn.click(
+        fn=analyze_image,
+        inputs=input_image,
+        outputs=[label_output, heatmap_output, info_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch()