Update app.py

app.py

@@ -15,18 +15,17 @@ MODEL_URL = (
     "https://huggingface.co/neuralninja10/deepFakeWithCBAM/"
     "resolve/main/deepFakeWithCBAM.pt"
 )
-
 MODEL_PATH = "deepFakeWithCBAM.pt"
 THRESHOLD = 0.68
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ============================================================
-# Page Config
+# Page Configuration
 # ============================================================
 
 st.set_page_config(
-    page_title="DeepFake Detection Demo",
-    page_icon="🧠",
+    page_title="DeepFake Detection",
+    page_icon="🛡️",
     layout="centered",
 )
 
@@ -34,16 +33,16 @@ st.set_page_config(
 # Secure Model Loader
 # ============================================================
 
-@st.cache_resource(show_spinner=False)
+@st.cache_resource
 def load_model():
     token = os.environ.get("HF_TOKEN")
     if token is None:
-        raise RuntimeError("HF_TOKEN is not set in Space secrets")
+        raise RuntimeError("HF_TOKEN not found in Space secrets")
 
     headers = {"Authorization": f"Bearer {token}"}
 
     if not os.path.exists(MODEL_PATH):
-        with st.spinner("Initializing deepfake detection model..."):
+        with st.spinner("Initializing system..."):
             response = requests.get(
                 MODEL_URL,
                 headers=headers,
@@ -52,7 +51,7 @@ def load_model():
             )
             response.raise_for_status()
             with open(MODEL_PATH, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
+                for chunk in response.iter_content(8192):
                     f.write(chunk)
 
     model = torch.jit.load(MODEL_PATH, map_location=DEVICE)
@@ -60,7 +59,7 @@ def load_model():
     return model
 
 # ============================================================
-# Preprocessing
+# Image Processing
 # ============================================================
 
 _transform = transforms.Compose([
@@ -82,21 +81,19 @@ def preprocess_image(image: Image.Image) -> torch.Tensor:
 def run_inference(model, image: Image.Image):
     tensor = preprocess_image(image).to(DEVICE)
 
-    start = time.time()
+    start_time = time.time()
     with torch.no_grad():
         logits = model(tensor)
-        prob = torch.sigmoid(logits).item()
-    latency_ms = (time.time() - start) * 1000
+        probability = torch.sigmoid(logits).item()
+    latency = (time.time() - start_time) * 1000
 
-    is_real = prob > THRESHOLD
-    confidence = prob if is_real else (1 - prob)
+    is_real = probability > THRESHOLD
+    confidence = probability if is_real else (1 - probability)
 
     return {
-        "prediction": "Real" if is_real else "Fake",
+        "label": "Real" if is_real else "Fake",
         "confidence": confidence,
-        "real_prob": prob,
-        "fake_prob": 1 - prob,
-        "latency_ms": latency_ms,
+        "latency": latency,
     }
 
 # ============================================================
@@ -104,80 +101,50 @@ def run_inference(model, image: Image.Image):
 # ============================================================
 
 def main():
-    st.title("🧠 DeepFake Detection")
-    st.markdown(
-        """
-        Upload a facial image to determine whether it is **Real** or **AI-Generated**.
-        This demo runs entirely on **CPU** using a TorchScript model.
-        """
-    )
+    st.title("DeepFake Detection")
+    st.caption("Upload an image to verify authenticity.")
 
     try:
         model = load_model()
     except Exception as e:
-        st.error("❌ Failed to load the model.")
+        st.error("System initialization failed.")
         st.exception(e)
         return
 
     uploaded_file = st.file_uploader(
-        "Upload an image (JPG / PNG)",
+        "Upload Image",
         type=["jpg", "jpeg", "png"],
-        accept_multiple_files=False,
     )
 
     if uploaded_file:
         image = Image.open(uploaded_file).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_container_width=True)
+        st.image(image, caption="Uploaded Image")
 
-        if st.button("Run DeepFake Detection"):
-            with st.spinner("Running inference..."):
-                result = run_inference(model, image)
+        st.divider()
 
-            st.divider()
-            st.subheader("Detection Result")
+        if st.button("Analyze Image"):
+            with st.spinner("Analyzing..."):
+                result = run_inference(model, image)
 
-            if result["prediction"] == "Real":
-                st.success("✅ Real Face Detected")
+            if result["label"] == "Real":
+                st.success("✔ Image appears to be authentic")
             else:
-                st.error("❌ Deepfake Detected")
+                st.error("✖ Image is likely manipulated")
 
             st.metric(
-                label="Confidence Score",
+                label="Confidence",
                 value=f"{result['confidence']:.2%}",
             )
 
             st.caption(
-                f"Inference latency: {result['latency_ms']:.1f} ms (CPU)"
+                f"Processing time: {result['latency']:.0f} ms"
             )
 
-            with st.expander("Detailed Probabilities"):
-                st.write(f"Real Probability: {result['real_prob']:.4f}")
-                st.write(f"Fake Probability: {result['fake_prob']:.4f}")
-
     st.divider()
-
     st.caption(
-        """
-        ⚠️ **Disclaimer**  
-        This system is provided for research and demonstration purposes only.
-        Predictions may be incorrect and should not be used as the sole basis
-        for real-world decisions.
-        """
+        "This demo caters all the available generators including Style GAN and Diffusion model variants"
+        "For further inquiries please feel free to contact uzair.mughal@unikrew.com"
     )
 
-    with st.expander("Model Information"):
-        st.markdown(
-            """
-            - **Architecture:** EfficientNet + CBAM  
-            - **Input Resolution:** 256×256  
-            - **Runtime:** CPU (TorchScript)  
-            - **Threshold:** 0.68  
-            - **Known Limitations:**  
-              - Heavy compression artifacts  
-              - Extreme lighting conditions  
-              - Occluded or profile faces  
-            """
-        )
-
 if __name__ == "__main__":
     main()