Upload app.py.py

app.py.py

@@ -0,0 +1,565 @@
+"""
+FYP4 SPAM DETECTION API
+FastAPI application for email spam detection using DeBERTa and ViT models
+"""
+
+import os
+import re
+import io
+import torch
+import torch.nn as nn
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional
+import PyPDF2
+import pdfplumber
+from PIL import Image
+from transformers import (
+    DebertaV2Model,
+    DebertaV2Tokenizer,
+    ViTModel,
+    ViTImageProcessor
+)
+
+# ================================
+# CONFIGURATION
+# ================================
+
+class Config:
+    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    TEXT_MODEL = 'microsoft/deberta-v3-base'
+    IMAGE_MODEL = 'google/vit-base-patch16-224-in21k'
+    TEXT_HIDDEN_DIM = 768
+    IMAGE_HIDDEN_DIM = 768
+    FUSION_DIM = 512
+    NUM_CLASSES = 2
+    DROPOUT = 0.3
+    MAX_TEXT_LENGTH = 256
+    IMG_SIZE = 224
+
+config = Config()
+
+# ================================
+# MODEL ARCHITECTURES
+# ================================
+
+class DeBERTaTextEncoder(nn.Module):
+    def __init__(self, dropout=0.3):
+        super(DeBERTaTextEncoder, self).__init__()
+        self.deberta = DebertaV2Model.from_pretrained(config.TEXT_MODEL)
+        self.projection = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(config.TEXT_HIDDEN_DIM, config.FUSION_DIM),
+            nn.LayerNorm(config.FUSION_DIM),
+            nn.GELU()
+        )
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask)
+        pooled = outputs.last_hidden_state[:, 0, :]
+        return self.projection(pooled)
+
+
+class ViTImageEncoder(nn.Module):
+    def __init__(self, dropout=0.3):
+        super(ViTImageEncoder, self).__init__()
+        self.vit = ViTModel.from_pretrained(config.IMAGE_MODEL)
+        self.projection = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(config.IMAGE_HIDDEN_DIM, config.FUSION_DIM),
+            nn.LayerNorm(config.FUSION_DIM),
+            nn.GELU()
+        )
+
+    def forward(self, pixel_values):
+        outputs = self.vit(pixel_values=pixel_values, return_dict=True)
+        pooled = outputs.last_hidden_state[:, 0, :]
+        return self.projection(pooled)
+
+
+class CrossModalAttention(nn.Module):
+    def __init__(self, dim=512, num_heads=8, dropout=0.1):
+        super(CrossModalAttention, self).__init__()
+        self.cross_attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, text_features, image_features):
+        text_features = text_features.unsqueeze(1)
+        image_features = image_features.unsqueeze(1)
+        attn_output, _ = self.cross_attn(text_features, image_features, image_features)
+        fused = self.norm1(text_features + attn_output)
+        ffn_output = self.ffn(fused)
+        output = self.norm2(fused + ffn_output)
+        return output.squeeze(1)
+
+
+class TextSpamClassifier(nn.Module):
+    def __init__(self, dropout=0.3):
+        super(TextSpamClassifier, self).__init__()
+        self.text_encoder = DeBERTaTextEncoder(dropout)
+        self.classifier = nn.Sequential(
+            nn.Linear(config.FUSION_DIM, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, 128),
+            nn.LayerNorm(128),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, config.NUM_CLASSES)
+        )
+
+    def forward(self, input_ids, attention_mask):
+        features = self.text_encoder(input_ids, attention_mask)
+        return self.classifier(features)
+
+
+class ImageSpamClassifier(nn.Module):
+    def __init__(self, dropout=0.3):
+        super(ImageSpamClassifier, self).__init__()
+        self.image_encoder = ViTImageEncoder(dropout)
+        self.classifier = nn.Sequential(
+            nn.Linear(config.FUSION_DIM, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, 128),
+            nn.LayerNorm(128),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, config.NUM_CLASSES)
+        )
+
+    def forward(self, pixel_values):
+        features = self.image_encoder(pixel_values)
+        return self.classifier(features)
+
+
+class FusionSpamClassifier(nn.Module):
+    def __init__(self, dropout=0.3):
+        super(FusionSpamClassifier, self).__init__()
+        self.text_encoder = DeBERTaTextEncoder(dropout)
+        self.image_encoder = ViTImageEncoder(dropout)
+        self.cross_modal_fusion = CrossModalAttention(config.FUSION_DIM, num_heads=8, dropout=dropout)
+        self.classifier = nn.Sequential(
+            nn.Linear(config.FUSION_DIM, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, 128),
+            nn.LayerNorm(128),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, config.NUM_CLASSES)
+        )
+
+    def forward(self, input_ids=None, attention_mask=None, pixel_values=None):
+        if input_ids is not None and pixel_values is not None:
+            text_features = self.text_encoder(input_ids, attention_mask)
+            image_features = self.image_encoder(pixel_values)
+            fused_features = self.cross_modal_fusion(text_features, image_features)
+        elif input_ids is not None:
+            fused_features = self.text_encoder(input_ids, attention_mask)
+        elif pixel_values is not None:
+            fused_features = self.image_encoder(pixel_values)
+        else:
+            raise ValueError("At least one modality required")
+        return self.classifier(fused_features)
+
+
+# ================================
+# PDF EXTRACTION
+# ================================
+
+class PDFExtractor:
+    @staticmethod
+    def extract_text_from_pdf(pdf_bytes):
+        """Extract text from PDF bytes"""
+        email_data = {
+            'subject': '',
+            'sender': '',
+            'body': '',
+            'full_text': ''
+        }
+
+        try:
+            pdf_file = io.BytesIO(pdf_bytes)
+            with pdfplumber.open(pdf_file) as pdf:
+                full_text = ""
+                for page in pdf.pages:
+                    text = page.extract_text()
+                    if text:
+                        full_text += text + "\n"
+
+                email_data['full_text'] = full_text
+
+                patterns = {
+                    'subject': [r'Subject:\s*(.+)', r'SUBJECT:\s*(.+)'],
+                    'sender': [r'From:\s*(.+)', r'FROM:\s*(.+)']
+                }
+
+                for field, pattern_list in patterns.items():
+                    for pattern in pattern_list:
+                        match = re.search(pattern, full_text, re.IGNORECASE)
+                        if match:
+                            email_data[field] = match.group(1).strip()[:100]
+                            break
+
+                body_match = re.search(r'(?:Subject|Date|From|To):.+?\n\n(.+)', full_text, re.DOTALL | re.IGNORECASE)
+                if body_match:
+                    email_data['body'] = body_match.group(1).strip()
+                else:
+                    email_data['body'] = full_text
+
+                return email_data
+        except Exception as e:
+            try:
+                pdf_file = io.BytesIO(pdf_bytes)
+                pdf_reader = PyPDF2.PdfReader(pdf_file)
+                full_text = ""
+                for page in pdf_reader.pages:
+                    text = page.extract_text()
+                    if text:
+                        full_text += text + "\n"
+                email_data['full_text'] = full_text
+                email_data['body'] = full_text
+                return email_data
+            except Exception as e:
+                raise HTTPException(status_code=400, detail=f"Error extracting text from PDF: {str(e)}")
+
+    @staticmethod
+    def extract_images_from_pdf(pdf_bytes):
+        """Extract first image from PDF bytes"""
+        try:
+            pdf_file = io.BytesIO(pdf_bytes)
+            pdf_reader = PyPDF2.PdfReader(pdf_file)
+            
+            for page_num, page in enumerate(pdf_reader.pages):
+                if '/XObject' in page['/Resources']:
+                    xObject = page['/Resources']['/XObject'].get_object()
+                    for obj in xObject:
+                        if xObject[obj]['/Subtype'] == '/Image':
+                            try:
+                                size = (xObject[obj]['/Width'], xObject[obj]['/Height'])
+                                data = xObject[obj].get_data()
+                                mode = "RGB" if xObject[obj]['/ColorSpace'] == '/DeviceRGB' else "P"
+                                img = Image.frombytes(mode, size, data)
+                                return img
+                            except:
+                                continue
+        except:
+            pass
+        return None
+
+
+# ================================
+# TEXT PREPROCESSING
+# ================================
+
+def preprocess_text(text):
+    """Preprocess text for model input"""
+    text = str(text).lower()
+    text = re.sub(r'http\S+|www\.\S+', '[URL]', text)
+    text = re.sub(r'\S+@\S+', '[EMAIL]', text)
+    text = re.sub(r'\d+', '[NUM]', text)
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text
+
+
+# ================================
+# SPAM DETECTOR
+# ================================
+
+class SpamDetector:
+    def __init__(self, text_model_path=None, image_model_path=None, fusion_model_path=None):
+        self.device = config.DEVICE
+        self.tokenizer = DebertaV2Tokenizer.from_pretrained(config.TEXT_MODEL)
+        self.image_processor = ViTImageProcessor.from_pretrained(config.IMAGE_MODEL)
+
+        self.text_model = None
+        self.image_model = None
+        self.fusion_model = None
+
+        # Load models
+        if text_model_path and os.path.exists(text_model_path):
+            print(f"Loading text model from {text_model_path}...")
+            self.text_model = TextSpamClassifier().to(self.device)
+            checkpoint = torch.load(text_model_path, map_location=self.device)
+            self.text_model.load_state_dict(checkpoint['model_state_dict'])
+            self.text_model.eval()
+            print("Text model loaded successfully!")
+
+        if image_model_path and os.path.exists(image_model_path):
+            print(f"Loading image model from {image_model_path}...")
+            self.image_model = ImageSpamClassifier().to(self.device)
+            checkpoint = torch.load(image_model_path, map_location=self.device)
+            self.image_model.load_state_dict(checkpoint['model_state_dict'])
+            self.image_model.eval()
+            print("Image model loaded successfully!")
+
+        if fusion_model_path and os.path.exists(fusion_model_path):
+            print(f"Loading fusion model from {fusion_model_path}...")
+            self.fusion_model = FusionSpamClassifier().to(self.device)
+            checkpoint = torch.load(fusion_model_path, map_location=self.device)
+            self.fusion_model.load_state_dict(checkpoint['model_state_dict'])
+            self.fusion_model.eval()
+            print("Fusion model loaded successfully!")
+
+    def predict_text(self, text):
+        if not self.text_model:
+            return None
+
+        encoding = self.tokenizer(
+            preprocess_text(text),
+            add_special_tokens=True,
+            max_length=config.MAX_TEXT_LENGTH,
+            padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors='pt'
+        )
+
+        input_ids = encoding['input_ids'].to(self.device)
+        attention_mask = encoding['attention_mask'].to(self.device)
+
+        with torch.no_grad():
+            outputs = self.text_model(input_ids, attention_mask)
+            probs = torch.softmax(outputs, dim=1)
+            predicted = torch.argmax(probs, dim=1)
+
+        return {
+            'prediction': 'SPAM' if predicted.item() == 1 else 'LEGITIMATE',
+            'confidence': float(probs[0, predicted.item()].item() * 100),
+            'spam_probability': float(probs[0, 1].item() * 100),
+            'ham_probability': float(probs[0, 0].item() * 100)
+        }
+
+    def predict_image(self, image):
+        if not self.image_model or image is None:
+            return None
+
+        try:
+            inputs = self.image_processor(images=image, return_tensors='pt')
+            pixel_values = inputs['pixel_values'].to(self.device)
+
+            with torch.no_grad():
+                outputs = self.image_model(pixel_values)
+                probs = torch.softmax(outputs, dim=1)
+                predicted = torch.argmax(probs, dim=1)
+
+            return {
+                'prediction': 'SPAM' if predicted.item() == 1 else 'LEGITIMATE',
+                'confidence': float(probs[0, predicted.item()].item() * 100),
+                'spam_probability': float(probs[0, 1].item() * 100),
+                'ham_probability': float(probs[0, 0].item() * 100)
+            }
+        except Exception as e:
+            return {'error': str(e)}
+
+    def predict_fusion(self, text, image=None):
+        if not self.fusion_model:
+            return None
+
+        encoding = self.tokenizer(
+            preprocess_text(text),
+            add_special_tokens=True,
+            max_length=config.MAX_TEXT_LENGTH,
+            padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors='pt'
+        )
+
+        input_ids = encoding['input_ids'].to(self.device)
+        attention_mask = encoding['attention_mask'].to(self.device)
+
+        pixel_values = None
+        if image is not None:
+            try:
+                image_inputs = self.image_processor(images=image, return_tensors='pt')
+                pixel_values = image_inputs['pixel_values'].to(self.device)
+            except:
+                pass
+
+        with torch.no_grad():
+            outputs = self.fusion_model(input_ids, attention_mask, pixel_values)
+            probs = torch.softmax(outputs, dim=1)
+            predicted = torch.argmax(probs, dim=1)
+
+        return {
+            'prediction': 'SPAM' if predicted.item() == 1 else 'LEGITIMATE',
+            'confidence': float(probs[0, predicted.item()].item() * 100),
+            'spam_probability': float(probs[0, 1].item() * 100),
+            'ham_probability': float(probs[0, 0].item() * 100)
+        }
+
+
+# ================================
+# FASTAPI APPLICATION
+# ================================
+
+app = FastAPI(
+    title="FYP4 Spam Detection API",
+    description="Email spam detection using DeBERTa and ViT models",
+    version="1.0.0"
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize detector (models will be loaded on startup)
+detector = None
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load models on startup"""
+    global detector
+    
+    text_model_path = os.getenv("TEXT_MODEL_PATH", "models/text_model.pth")
+    image_model_path = os.getenv("IMAGE_MODEL_PATH", "models/image_model.pth")
+    fusion_model_path = os.getenv("FUSION_MODEL_PATH", "models/fusion_model.pth")
+    
+    # Check which models exist
+    text_exists = os.path.exists(text_model_path)
+    image_exists = os.path.exists(image_model_path)
+    fusion_exists = os.path.exists(fusion_model_path)
+    
+    print(f"Models availability: Text={text_exists}, Image={image_exists}, Fusion={fusion_exists}")
+    
+    detector = SpamDetector(
+        text_model_path=text_model_path if text_exists else None,
+        image_model_path=image_model_path if image_exists else None,
+        fusion_model_path=fusion_model_path if fusion_exists else None
+    )
+    
+    print("API ready!")
+
+
+# Pydantic models for request/response
+class TextRequest(BaseModel):
+    text: str
+
+
+class PredictionResponse(BaseModel):
+    prediction: str
+    confidence: float
+    spam_probability: float
+    ham_probability: float
+    model_used: str
+
+
+class PDFPredictionResponse(BaseModel):
+    email_data: dict
+    text_result: Optional[dict]
+    image_result: Optional[dict]
+    fusion_result: Optional[dict]
+    final_prediction: str
+    final_confidence: float
+
+
+@app.get("/")
+async def root():
+    """Root endpoint with API information"""
+    return {
+        "message": "FYP4 Spam Detection API",
+        "version": "1.0.0",
+        "endpoints": {
+            "POST /predict/text": "Predict spam from text",
+            "POST /predict/pdf": "Predict spam from PDF email",
+            "GET /health": "Health check"
+        }
+    }
+
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "device": str(config.DEVICE),
+        "models_loaded": {
+            "text": detector.text_model is not None if detector else False,
+            "image": detector.image_model is not None if detector else False,
+            "fusion": detector.fusion_model is not None if detector else False
+        }
+    }
+
+
+@app.post("/predict/text", response_model=PredictionResponse)
+async def predict_text(request: TextRequest):
+    """Predict spam from text content"""
+    if not detector or not detector.text_model:
+        raise HTTPException(status_code=503, detail="Text model not available")
+    
+    result = detector.predict_text(request.text)
+    result['model_used'] = 'text'
+    
+    return result
+
+
+@app.post("/predict/pdf", response_model=PDFPredictionResponse)
+async def predict_pdf(file: UploadFile = File(...)):
+    """Predict spam from PDF email"""
+    if not file.filename.endswith('.pdf'):
+        raise HTTPException(status_code=400, detail="File must be a PDF")
+    
+    if not detector:
+        raise HTTPException(status_code=503, detail="Models not loaded")
+    
+    # Read PDF
+    pdf_bytes = await file.read()
+    
+    # Extract text and images
+    email_data = PDFExtractor.extract_text_from_pdf(pdf_bytes)
+    full_text = f"{email_data['subject']}\n\n{email_data['body']}"
+    image = PDFExtractor.extract_images_from_pdf(pdf_bytes)
+    
+    # Get predictions
+    results = {
+        'email_data': email_data,
+        'text_result': None,
+        'image_result': None,
+        'fusion_result': None
+    }
+    
+    if detector.text_model:
+        results['text_result'] = detector.predict_text(full_text)
+    
+    if detector.image_model and image:
+        results['image_result'] = detector.predict_image(image)
+    
+    if detector.fusion_model:
+        results['fusion_result'] = detector.predict_fusion(full_text, image)
+    
+    # Determine final prediction (prioritize: fusion > text > image)
+    final_result = results['fusion_result'] or results['text_result'] or results['image_result']
+    
+    if not final_result:
+        raise HTTPException(status_code=503, detail="No models available for prediction")
+    
+    results['final_prediction'] = final_result['prediction']
+    results['final_confidence'] = final_result['confidence']
+    
+    return results
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)