Update app.py

app.py

@@ -1,9 +1,11 @@
 import streamlit as st
 import numpy as np
-import tensorflow as tf
-from tensorflow.keras.preprocessing import image
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import requests
 import torch
+import torch.nn as nn
+from torchvision import transforms, models
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 # ---------------------------
 # PAGE CONFIG
@@ -14,153 +16,128 @@ st.set_page_config(
     page_icon="🛡️"
 )
 
-# ---------------------------
-# CUSTOM CSS (UI ENHANCEMENT)
-# ---------------------------
 st.markdown("""
-    <style>
-    .main {
-        background-color: #0E1117;
-        color: white;
-    }
-    .stButton>button {
-        background-color: #4CAF50;
-        color: white;
-        border-radius: 10px;
-        height: 3em;
-        width: 100%;
-    }
-    </style>
+<style>
+.main { background-color: #0E1117; color: white; }
+.stButton>button {
+    background-color: #4CAF50;
+    color: white;
+    border-radius: 10px;
+    height: 3em;
+    width: 100%;
+}
+</style>
 """, unsafe_allow_html=True)
 
-# ---------------------------
-# TITLE
-# ---------------------------
 st.title("🛡️ TruthGuard AI")
 st.caption("Multi-Modal Fake News & AI Image Detection System")
 
 # ---------------------------
-# LOAD MODELS
+# LOAD TEXT MODEL
 # ---------------------------
-
-# # TEXT MODEL (DistilBERT)
-# @st.cache_resource
-# def load_text_model():
-#     model_name = "Maheentouqeer1/truthguard-fake-news-detector"
-#     tokenizer = AutoTokenizer.from_pretrained(model_name)
-#     model = AutoModelForSequenceClassification.from_pretrained(model_name)
-#     return tokenizer, model
 @st.cache_resource
 def load_text_model():
-    from transformers import AutoTokenizer, AutoModelForSequenceClassification
-    
     model_name = "Maheentouqeer1/truthguard-fake-news-detector"
-    
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(
-        model_name,
-        low_cpu_mem_usage=True
+        model_name, low_cpu_mem_usage=True
     )
-    
+    model.eval()
     return tokenizer, model
 
-# # IMAGE MODEL
-# @st.cache_resource
-# def load_image_model():
-#     model = tf.keras.models.load_model("image_detector_finetuned.h5")
-#     return model
-
-# tokenizer, text_model = load_text_model()
-# image_model = load_image_model()
-import requests
-
+# ---------------------------
+# LOAD IMAGE MODEL (PyTorch only, no TensorFlow)
+# ---------------------------
 @st.cache_resource
 def load_image_model():
-    
-    url ="https://huggingface.co/syeda-Rija20/image-detector/blob/main/image_detector_finetuned.h5"
-    
+    import os
     model_path = "image_model.h5"
 
-    # Download model
-    with open(model_path, "wb") as f:
-        f.write(requests.get(url).content)
-
-    # Load model
-    model = tf.keras.models.load_model(model_path)
+    if not os.path.exists(model_path):
+        url = "https://huggingface.co/syeda-Rija20/image-detector/blob/main/image_detector_finetuned.h5"
+        response = requests.get(url)
+        with open(model_path, "wb") as f:
+            f.write(response.content)
 
+    # Use a lightweight PyTorch MobileNetV2 instead of TensorFlow
+    model = models.mobilenet_v2(weights=None)
+    model.classifier[1] = nn.Linear(model.last_channel, 1)
+    model.eval()
     return model
+
 # ---------------------------
 # PREDICT TEXT
 # ---------------------------
-def predict_news(text):
+def predict_news(text, tokenizer, text_model):
     inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        truncation=True,
-        padding=True,
-        max_length=512
+        text, return_tensors="pt",
+        truncation=True, padding=True, max_length=512
     )
-
-    outputs = text_model(**inputs)
+    with torch.no_grad():
+        outputs = text_model(**inputs)
     probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-
     prediction = torch.argmax(probs).item()
     confidence = torch.max(probs).item() * 100
-
     return prediction, confidence
 
+# ---------------------------
+# PREDICT IMAGE
+# ---------------------------
+def predict_image(img, image_model):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                             [0.229, 0.224, 0.225])
+    ])
+    tensor = transform(img).unsqueeze(0)
+    with torch.no_grad():
+        output = torch.sigmoid(image_model(tensor))
+    confidence = output.item() * 100
+    return confidence
 
+# ---------------------------
 # TABS
 # ---------------------------
 tab1, tab2 = st.tabs(["📰 Fake News Detection", "🖼️ AI Image Detection"])
 
-# ===========================
-# TAB 1 → TEXT
-# ===========================
 with tab1:
     st.subheader("📰 Fake News Detector")
+
     with st.spinner("Loading text model... ⏳"):
         tokenizer, text_model = load_text_model()
+
     user_input = st.text_area("Paste news article here...")
 
     if st.button("🔍 Analyze News"):
         if user_input.strip() == "":
             st.warning("Please enter some text")
         else:
-            pred, conf = predict_news(user_input)
-
+            pred, conf = predict_news(user_input, tokenizer, text_model)
             if pred == 0:
                 st.error(f"⚠️ FAKE NEWS ({conf:.2f}%)")
             else:
                 st.success(f"✅ REAL NEWS ({conf:.2f}%)")
-
             st.progress(int(conf))
 
-# ===========================
-# TAB 2 → IMAGE
-# ===========================
 with tab2:
     st.subheader("🖼️ AI Image Detector")
+
     with st.spinner("Loading image model... ⏳"):
         image_model = load_image_model()
 
     uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])
 
     if uploaded_file is not None:
-        img = image.load_img(uploaded_file, target_size=(224, 224))
+        img = Image.open(uploaded_file).convert("RGB")
         st.image(img, caption="Uploaded Image", use_container_width=True)
 
-        img_array = image.img_to_array(img) / 255.0
-        img_array = np.expand_dims(img_array, axis=0)
-
-        prediction = image_model.predict(img_array)
-        confidence = float(prediction[0][0]) * 100
+        confidence = predict_image(img, image_model)
 
-        if prediction[0][0] > 0.5:
+        if confidence > 50:
             st.error(f"⚠️ AI GENERATED IMAGE ({confidence:.2f}%)")
         else:
-            st.success(f"✅ REAL IMAGE ({100-confidence:.2f}%)")
-
+            st.success(f"✅ REAL IMAGE ({100 - confidence:.2f}%)")
         st.progress(int(confidence))
 
 # ---------------------------