Max005/DeepfakeDetection

1from fastapi import FastAPI, File, UploadFile2from pydantic import BaseModel3import uvicorn4import os5import torchaudio6import torch.nn.functional as F7import torch8from transformers import AutoProcessor, AutoModelForAudioClassification9from pathlib import Path10 11app_dir = Path(__file__).parent12# Model setup13model_path = app_dir / "Deepfake" / "model"14 15processor = AutoProcessor.from_pretrained(model_path)16model = AutoModelForAudioClassification.from_pretrained(17    pretrained_model_name_or_path=model_path,18    local_files_only=True,19)20 21def prepare_audio(file_path, sampling_rate=16000, duration=10):22    """23    Prepares audio by loading, resampling, and returning it in manageable chunks.24    """25    # Load and resample the audio file26    waveform, original_sampling_rate = torchaudio.load(file_path)27    28    # Convert stereo to mono if necessary29    if waveform.shape[0] > 1:  # More than 1 channel30        waveform = torch.mean(waveform, dim=0, keepdim=True)31    32    # Resample if needed33    if original_sampling_rate != sampling_rate:34        resampler = torchaudio.transforms.Resample(orig_freq=original_sampling_rate, new_freq=sampling_rate)35        waveform = resampler(waveform)36    37    # Calculate chunk size in samples38    chunk_size = sampling_rate * duration39    audio_chunks = []40 41    # Split the audio into chunks42    for start in range(0, waveform.shape[1], chunk_size):43        chunk = waveform[:, start:start + chunk_size]44        45        # Pad the last chunk if it's shorter than the chunk size46        if chunk.shape[1] < chunk_size:47            padding = chunk_size - chunk.shape[1]48            chunk = torch.nn.functional.pad(chunk, (0, padding))49        50        audio_chunks.append(chunk.squeeze().numpy())51    52    return audio_chunks53 54def predict_audio(file_path):55    """56    Predicts the class of an audio file by aggregating predictions from chunks and calculates confidence.57    """58    # Prepare audio chunks59    audio_chunks = prepare_audio(file_path)60    predictions = []61    confidences = []62 63    for i, chunk in enumerate(audio_chunks):64        # Prepare input for the model65        inputs = processor(66            chunk, sampling_rate=16000, return_tensors="pt", padding=True67        )68        69        # Perform inference70        with torch.no_grad():71            outputs = model(**inputs)72            logits = outputs.logits73            74            # Apply softmax to calculate probabilities75            probabilities = F.softmax(logits, dim=1)76            77            # Get the predicted class and its confidence78            confidence, predicted_class = torch.max(probabilities, dim=1)79            predictions.append(predicted_class.item())80            confidences.append(confidence.item())81    82    # Aggregate predictions (majority voting)83    aggregated_prediction_id = max(set(predictions), key=predictions.count)84    predicted_label = model.config.id2label[aggregated_prediction_id]85    86    # Calculate average confidence across chunks87    average_confidence = sum(confidences) / len(confidences)88 89    return {90        "predicted_label": predicted_label,91        "average_confidence": average_confidence92    }93 94# Initialize FastAPI95app = FastAPI()96 97@app.post("/infer")98async def infer(file: UploadFile = File(...)):99    """100    Accepts an audio file and returns the prediction and confidence.101    """102    # Save the uploaded file to a temporary location103    temp_file_path = f"temp_{file.filename}"104    with open(temp_file_path, "wb") as temp_file:105        temp_file.write(await file.read())106    107    try:108        # Perform inference109        predictions = predict_audio(temp_file_path)110    finally:111        # Clean up the temporary file112        os.remove(temp_file_path)113    114    return predictions115 116@app.get("/health")117async def health():118    return {119        "message": "ok",120        "Sound":str(torchaudio.list_audio_backends())121        }122 123