Update app.py

app.py

@@ -1,72 +1,119 @@
 # app.py
-import torch
-from torchvision import transforms, models
-from PIL import Image
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, logging
+logging.set_verbosity_error()  # mute transformers INFO logs to keep the UI logs clean
 
-# -----------------------------
-# CONFIG
-# -----------------------------
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-IMAGE_MODEL_PATH = "models/deepfake_model.pth"
+# -----------------------
+# NOTE: These are public models known to load on Spaces.
+# - text model: small DistilBERT sentiment model (used as a safe demo for text "credibility")
+# - image model: ViT image-classifier (generic). Replace later with a custom deepfake model when ready.
+# -----------------------
+TEXT_MODEL_ID = "distilbert-base-uncased-finetuned-sst-2-english"
+IMAGE_MODEL_ID = "google/vit-base-patch16-224"
 
-# -----------------------------
-# LOAD DEEPFAKE IMAGE MODEL
-# -----------------------------
-deepfake_model = models.resnet18(pretrained=True)
-num_features = deepfake_model.fc.in_features
-deepfake_model.fc = torch.nn.Linear(num_features, 2)
-deepfake_model.load_state_dict(torch.load(IMAGE_MODEL_PATH, map_location=DEVICE))
-deepfake_model.eval()
-deepfake_model.to(DEVICE)
+# Load pipelines (will download weights on first run)
+try:
+    text_pipe = pipeline("text-classification", model=TEXT_MODEL_ID)
+except Exception as e:
+    text_pipe = None
+    text_load_error = str(e)
+else:
+    text_load_error = None
 
-# -----------------------------
-# LOAD TEXT MODEL (already working)
-# -----------------------------
-text_model = pipeline("text-classification", model="mrm8488/bert-tiny-finetuned-fake-news")
+try:
+    image_pipe = pipeline("image-classification", model=IMAGE_MODEL_ID)
+except Exception as e:
+    image_pipe = None
+    image_load_error = str(e)
+else:
+    image_load_error = None
 
-# -----------------------------
-# IMAGE PREDICTION FUNCTION
-# -----------------------------
-image_transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-])
+# Friendly mapping (different text models return different label names)
+TEXT_FRIENDLY = {
+    "NEGATIVE": "Not credible / Fake (demo)",
+    "LABEL_0": "Not credible / Fake (demo)",
+    "POSITIVE": "Credible / Real (demo)",
+    "LABEL_1": "Credible / Real (demo)"
+}
 
-def predict_image(image):
-    img = Image.fromarray(image)
-    img = image_transform(img).unsqueeze(0).to(DEVICE)
-    with torch.no_grad():
-        output = deepfake_model(img)
-        probs = torch.softmax(output, dim=1)
-        label = "Real" if torch.argmax(probs) == 0 else "Fake"
-        confidence = torch.max(probs).item()
-    return f"{label} ({confidence*100:.2f}% confidence)"
+def friendly_text_label(raw_label: str) -> str:
+    if raw_label is None:
+        return "Unknown"
+    raw = str(raw_label).upper()
+    return TEXT_FRIENDLY.get(raw, raw_label)
 
-# -----------------------------
-# TEXT PREDICTION FUNCTION
-# -----------------------------
-def predict_text(text):
-    result = text_model(text)[0]
-    return f"{result['label']} ({result['score']*100:.2f}% confidence)"
+def classify_text(text: str):
+    if text is None or not str(text).strip():
+        return "Please paste some text to analyze.", {}
+    if text_pipe is None:
+        return f"Text model failed to load: {text_load_error}", {}
+    try:
+        # request top 2 classes for a small confidence breakdown
+        preds = text_pipe(text, top_k=2)
+        # preds is a list of dicts like {'label':'POSITIVE','score':0.98}
+        label_dict = {}
+        for p in preds:
+            lab = friendly_text_label(p.get("label"))
+            label_dict[lab] = float(p.get("score", 0.0))
+        # choose top
+        top_lab = max(label_dict.items(), key=lambda kv: kv[1])
+        summary = f"{top_lab[0]} ({top_lab[1]*100:.2f}%)"
+        return summary, label_dict
+    except Exception as e:
+        return f"Error during text classification: {e}", {}
 
-# -----------------------------
-# BUILD UI
-# -----------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## Deepfake & Fake News Detector")
-    
-    with gr.Tab("Image Detection"):
-        image_input = gr.Image(type="numpy")
-        image_output = gr.Textbox(label="Prediction")
-        image_button = gr.Button("Predict Image")
-        image_button.click(predict_image, inputs=image_input, outputs=image_output)
-    
-    with gr.Tab("Text Detection"):
-        text_input = gr.Textbox(label="Enter Text")
-        text_output = gr.Textbox(label="Prediction")
-        text_button = gr.Button("Predict Text")
-        text_button.click(predict_text, inputs=text_input, outputs=text_output)
+def classify_image(image):
+    if image is None:
+        return "Please upload an image.", {}
+    if image_pipe is None:
+        return f"Image model failed to load: {image_load_error}", {}
+    try:
+        preds = image_pipe(image, top_k=5)
+        label_dict = {p['label']: float(p['score']) for p in preds}
+        top_lab = max(label_dict.items(), key=lambda kv: kv[1])
+        summary = f"{top_lab[0]} ({top_lab[1]*100:.2f}%)"
+        return summary, label_dict
+    except Exception as e:
+        return f"Error during image classification: {e}", {}
+
+# --- UI ---
+with gr.Blocks(title="AI Detector (Text + Image)") as demo:
+    gr.Markdown("## 🔎 AI Detector\nText (credibility demo) and Image (generic classifier).")
+    gr.Markdown(
+        "> This app uses public models that load in Spaces. When you have your own trained deepfake model, "
+        "you can swap the image model ID in `app.py` to point at your Hugging Face model."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📝 Text Analysis")
+            txt = gr.Textbox(lines=6, placeholder="Paste text here...", label="Input Text")
+            txt_result = gr.Textbox(label="Summary")
+            txt_probs = gr.Label(label="Confidence (top 2)")
+            with gr.Row():
+                btn_txt = gr.Button("Analyze Text")
+                btn_txt.click(classify_text, inputs=txt, outputs=[txt_result, txt_probs])
+                btn_txt_clear = gr.Button("Clear")
+                btn_txt_clear.click(lambda: ("", {}, ""), outputs=[txt, txt_result, txt_probs])
+
+        with gr.Column(scale=1):
+            gr.Markdown("### 🖼️ Image Analysis")
+            img = gr.Image(type="pil", label="Upload Image")
+            img_result = gr.Textbox(label="Summary")
+            img_probs = gr.Label(label="Top-5 Confidence")
+            with gr.Row():
+                btn_img = gr.Button("Analyze Image")
+                btn_img.click(classify_image, inputs=img, outputs=[img_result, img_probs])
+                btn_img_clear = gr.Button("Clear")
+                btn_img_clear.click(lambda: (None, "", {}), outputs=[img, img_result, img_probs])
+
+    # Footer: show load errors if any
+    if text_load_error or image_load_error:
+        with gr.Column():
+            gr.Markdown("**Model load warnings:**")
+            if text_load_error:
+                gr.Markdown(f"- Text model load error: `{text_load_error}`")
+            if image_load_error:
+                gr.Markdown(f"- Image model load error: `{image_load_error}`")
 
 demo.launch()