Update README.md

README.md

@@ -86,65 +86,153 @@ Final metrics on validation set:
 - Precision: 0.9692 (96.92%)
 - Recall: 0.9728 (97.28%)
 
-## Usage
-
-Here's how to use the model:
+# Usage Guide
 
+## 1. Environment Setup
+First, clone the AST repository and install required dependencies:
 ```python
+# Clone AST repository and set up path
+git clone https://github.com/YuanGongND/ast.git
+import sys
+sys.path.append('./ast')
+cd ast
+
+# Install dependencies
+pip install timm==0.4.5 wget
+
+# Required imports
+import os
 import torch
 import torchaudio
-from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
-
-# Load model and processor 
-model = AutoModelForAudioClassification.from_pretrained("WpythonW/ast-fakeaudio-detector")
-processor = AutoFeatureExtractor.from_pretrained("WpythonW/ast-fakeaudio-detector")
-
-# Load audio file
-audio_path = "/content/r.mp3" 
-waveform, sample_rate = torchaudio.load(audio_path)
-
-# Resample to required sample rate
-if sample_rate != 16000:
-    resampler = torchaudio.transforms.Resample(sample_rate, 16000)
-    waveform = resampler(waveform)
-
-# Convert stereo to mono if needed
-if waveform.shape[0] > 1:
-    waveform = torch.mean(waveform, dim=0, keepdim=True)
-
-# Create mel spectrogram
-mel_spec = torchaudio.transforms.MelSpectrogram(
-    sample_rate=16000,
-    n_mels=128,
-    n_fft=2048, 
-    hop_length=160
-)(waveform)
-
-# Convert to decibels
-mel_spec_db = torchaudio.transforms.AmplitudeToDB()(mel_spec)
-
-# Normalize
-mel_spec_norm = (mel_spec_db + 4.26) / (4.57 * 2)
-
-# Adjust to required length (1024 frames)
-if mel_spec_norm.shape[2] < 1024:
-    # If shorter - pad with zeros
-    padding = torch.zeros(1, 128, 1024 - mel_spec_norm.shape[2])
-    mel_spec_norm = torch.cat([mel_spec_norm, padding], dim=2)
-else:
-    # If longer - truncate
-    mel_spec_norm = mel_spec_norm[:, :, :1024]
-
-# Get prediction  
+import matplotlib.pyplot as plt
+import numpy as np
+from torch import nn
+from src.models import ASTModel
+```
+
+## 2. Model Implementation
+Implement the BinaryAST model class:
+```python
+class BinaryAST(nn.Module):
+    def __init__(self, pretrained_path='pretrained_models/audioset_10_10_0.4593.pth'):
+        super().__init__()
+        # Initialize AST base model
+        self.ast = ASTModel(
+            label_dim=527,
+            input_fdim=128,
+            input_tdim=1024,
+            imagenet_pretrain=True,
+            audioset_pretrain=False,
+            model_size='base384'
+        )
+        
+        # Load pretrained weights if available
+        if os.path.exists(pretrained_path):
+            print(f"Loading pretrained weights from {pretrained_path}")
+            state_dict = torch.load(pretrained_path, map_location='cpu', weights_only=True)
+            self.ast.load_state_dict(state_dict, strict=False)
+
+        # Binary classification head
+        self.ast.mlp_head = nn.Sequential(
+            nn.LayerNorm(768),
+            nn.Dropout(0.3),
+            nn.Linear(768, 1)
+        )
+
+    def forward(self, x):
+        return self.ast(x)
+```
+
+## 3. Audio Processing Function
+Function to preprocess audio files for model input:
+```python
+def process_audio(file_path, sr=16000):
+    """
+    Process audio file for model inference.
+    
+    Args:
+        file_path (str): Path to audio file
+        sr (int): Target sample rate (default: 16000)
+    
+    Returns:
+        torch.Tensor: Processed mel spectrogram (1024 x 128)
+    """
+    # Load audio
+    audio_tensor, orig_sr = torchaudio.load(file_path)
+    print(f"Initial tensor shape: {audio_tensor.shape}, sample_rate={orig_sr}")
+
+    # Convert to mono if needed
+    if audio_tensor.shape[0] > 1:
+        audio_tensor = torch.mean(audio_tensor, dim=0, keepdim=True)
+
+    # Resample to target sample rate
+    if orig_sr != sr:
+        resampler = torchaudio.transforms.Resample(orig_sr, sr)
+        audio_tensor = resampler(audio_tensor)
+
+    # Create mel spectrogram
+    mel_spec = torchaudio.transforms.MelSpectrogram(
+        sample_rate=sr,
+        n_mels=128,
+        n_fft=2048,
+        hop_length=160
+    )(audio_tensor)
+    spec_db = torchaudio.transforms.AmplitudeToDB()(mel_spec)
+
+    # Post-process spectrogram
+    spec_db = spec_db.squeeze(0).transpose(0, 1)
+    spec_db = (spec_db + 4.26) / (4.57 * 2)  # Normalize
+    
+    # Ensure correct length (pad/trim to 1024 frames)
+    target_len = 1024
+    if spec_db.shape[0] < target_len:
+        pad = torch.zeros(target_len - spec_db.shape[0], 128)
+        spec_db = torch.cat([spec_db, pad], dim=0)
+    else:
+        spec_db = spec_db[:target_len, :]
+
+    return spec_db
+```
+
+## 4. Model Loading and Inference
+Example of loading the model and running inference:
+```python
+# Initialize and load model
+model = BinaryAST()
+checkpoint = torch.load('/content/final_model.pth', map_location='cpu')
+model.load_state_dict(checkpoint['model_state_dict'])
+model.eval()
+
+# Process audio file
+spec = process_audio('path_to_audio.mp3')
+
+# Visualize spectrogram (optional)
+plt.figure(figsize=(10, 3))
+plt.imshow(spec.numpy().T, aspect='auto', origin='lower')
+plt.title('Mel Spectrogram')
+plt.xlabel('Time Frames')
+plt.ylabel('Mel Bins')
+plt.colorbar()
+plt.show()
+
+# Run inference
+spec_batch = spec.unsqueeze(0)
 with torch.no_grad():
-    outputs = model(mel_spec_norm)
-    probabilities = torch.sigmoid(outputs.logits)
-    is_fake = probabilities > 0.5
+    output = model(spec_batch)
+    prob_fake = torch.sigmoid(output).item()
 
-print(f"Probability of fake audio: {probabilities[0][0]:.4f}")
-print(f"Prediction: {'FAKE' if is_fake[0][0] else 'REAL'} audio")
+print(f"Probability of fake audio: {prob_fake:.4f}")
+print("Prediction:", "FAKE" if prob_fake > 0.5 else "REAL")
 ```
 
+## Key Notes:
+- Ensure audio files are accessible and in a supported format
+- The model expects 16kHz sample rate input
+- Input audio is converted to mono if stereo
+- The model outputs probability scores (>0.5 indicates fake audio)
+- Visualization of spectrograms is optional but useful for debugging
+
+
 ## Limitations
 
 Important considerations when using this model: