Max005/DeepfakeDetection

1from fastapi import FastAPI, File, UploadFile2from pydantic import BaseModel3import os4import torchaudio5import torch.nn.functional as F6import torch7from transformers import AutoProcessor, AutoModelForAudioClassification, pipeline8from pathlib import Path9 10app_dir = Path(__file__).parent11 12# Deepfake model setup13deepfake_model_path = app_dir / "Deepfake" / "model"14deepfake_processor = AutoProcessor.from_pretrained(deepfake_model_path)15deepfake_model = AutoModelForAudioClassification.from_pretrained(16    pretrained_model_name_or_path=deepfake_model_path,17    local_files_only=True,18)19 20def prepare_audio(file_path, sampling_rate=16000, duration=10):21    waveform, original_sampling_rate = torchaudio.load(file_path)22    if waveform.shape[0] > 1:23        waveform = torch.mean(waveform, dim=0, keepdim=True)24    if original_sampling_rate != sampling_rate:25        resampler = torchaudio.transforms.Resample(orig_freq=original_sampling_rate, new_freq=sampling_rate)26        waveform = resampler(waveform)27    chunk_size = sampling_rate * duration28    audio_chunks = []29    for start in range(0, waveform.shape[1], chunk_size):30        chunk = waveform[:, start:start + chunk_size]31        if chunk.shape[1] < chunk_size:32            padding = chunk_size - chunk.shape[1]33            chunk = torch.nn.functional.pad(chunk, (0, padding))34        audio_chunks.append(chunk.squeeze().numpy())35    return audio_chunks36 37def predict_audio(file_path):38    audio_chunks = prepare_audio(file_path)39    predictions = []40    confidences = []41    for chunk in audio_chunks:42        inputs = deepfake_processor(43            chunk, sampling_rate=16000, return_tensors="pt", padding=True44        )45        with torch.no_grad():46            outputs = deepfake_model(**inputs)47            logits = outputs.logits48            probabilities = F.softmax(logits, dim=1)49            confidence, predicted_class = torch.max(probabilities, dim=1)50            predictions.append(predicted_class.item())51            confidences.append(confidence.item())52    aggregated_prediction_id = max(set(predictions), key=predictions.count)53    predicted_label = deepfake_model.config.id2label[aggregated_prediction_id]54    average_confidence = sum(confidences) / len(confidences)55    return {56        "predicted_label": predicted_label,57        "average_confidence": average_confidence58    }59 60# ScamText model setup61scamtext_pipe = pipeline("text-classification", model="phishbot/ScamLLM")62 63# Input model for scam text inference64class TextInput(BaseModel):65    input: str66 67 68# Initialize FastAPI69app = FastAPI()70 71@app.post("/deepfake/infer")72async def deepfake_infer(file: UploadFile = File(...)):73    temp_file_path = f"temp_{file.filename}"74    with open(temp_file_path, "wb") as temp_file:75        temp_file.write(await file.read())76    try:77        predictions = predict_audio(temp_file_path)78    finally:79        os.remove(temp_file_path)80    return predictions81 82@app.post("/scamtext/infer")83async def scamtext_infer(data: TextInput):84    predictions = scamtext_pipe(data.input)85    return predictions86 87@app.get("/deepfake/health")88async def deepfake_health():89    return {90        "message": "ok",91        "Sound": str(torchaudio.list_audio_backends())92    }93 94@app.get("/scamtext/health")95async def scamtext_health():96    return {"message": "ok"}97 98if __name__ == "__main__":99    import uvicorn100    uvicorn.run(app, host="0.0.0.0", port=8000)101