predict.py

import os
import sys
import argparse
import logging
import tempfile
import wave
from pathlib import Path
from typing import Optional, Tuple, Any

import numpy as np
import librosa

from feature_extraction import AudioFeatureExtractor

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

try:
    import pyaudio
    PYAUDIO_AVAILABLE = True
except ImportError:
    PYAUDIO_AVAILABLE = False
    logger.debug("PyAudio not available - recording feature disabled")


class Config:
    MODELS_DIR = './models'
    BEST_MODEL_FILE = './best_model.txt'
    USE_PCA = None
    SAMPLE_RATE = 22050
    N_MFCC = 20
    N_CHROMA = 12
    CONFIDENCE_DECIMALS = 2


class ModelLoader:
    def __init__(self, models_dir: str = Config.MODELS_DIR):
        self.models_dir = Path(models_dir)
        self.scaler = None
        self.pca = None
        self.model = None
        self.model_name = None
        self.use_pca = False

    def get_best_model_name(self, best_model_file: str = Config.BEST_MODEL_FILE) -> str:
        if not os.path.exists(best_model_file):
            raise FileNotFoundError(
                f"Best model file not found: {best_model_file}\n"
                "Please run train.py and pca_train.py first."
            )

        with open(best_model_file, 'r') as f:
            model_name = f.read().strip()

        if not model_name:
            raise ValueError("Best model file is empty")

        return model_name

    def load_scaler(self, model_name: str) -> None:
        self.use_pca = '_PCA' in model_name or '_pca' in model_name.lower()

        if self.use_pca:
            scaler_path = self.models_dir / 'scaler_pca.joblib'
        else:
            scaler_path = self.models_dir / 'scaler_no_pca.joblib'

        if not scaler_path.exists():
            raise FileNotFoundError(f"Scaler not found: {scaler_path}")

        import joblib
        self.scaler = joblib.load(str(scaler_path))
       

    def load_pca(self) -> None:
        if not self.use_pca:
            return

        pca_path = self.models_dir / 'pca_model.joblib'
        if not pca_path.exists():
            raise FileNotFoundError(f"PCA model not found: {pca_path}")

        import joblib
        self.pca = joblib.load(str(pca_path))
       

    def load_model(self, model_name: str) -> None:
        clean_name = model_name.replace('_PCA', '').replace('_pca', '')
        suffix = '_pca.joblib' if self.use_pca else '_no_pca.joblib'

        model_path = self.models_dir / f'{clean_name}{suffix}'

        if not model_path.exists():
            alt_path = self.models_dir / f'{model_name}.joblib'
            if alt_path.exists():
                model_path = alt_path
            else:
                raise FileNotFoundError(
                    f"Model not found: {model_path}\n"
                    f"Available models in: {self.models_dir}"
                )

        import joblib
        self.model = joblib.load(str(model_path))
        self.model_name = model_name
        

    def load_all(self, model_name: Optional[str] = None) -> Tuple[Any, bool]:
        if model_name is None:
            model_name = self.get_best_model_name()

        self.model_name = model_name

  
        self.load_scaler(model_name)
        self.load_pca()
        self.load_model(model_name)

        return self.model, self.use_pca


class ParkinsonPredictor:
    def __init__(self, model_name: Optional[str] = None,
                 models_dir: str = Config.MODELS_DIR):
        self.models_dir = models_dir
        self.model_name = model_name
        self.model = None
        self.scaler = None
        self.pca = None
        self.use_pca = False
        self.feature_extractor = None
        self._initialized = False

    def initialize(self) -> None:
        loader = ModelLoader(self.models_dir)
        self.model, self.use_pca = loader.load_all(self.model_name)
        self.scaler = loader.scaler
        self.pca = loader.pca

        self.model_name = loader.model_name

        self.feature_extractor = AudioFeatureExtractor(
            sample_rate=Config.SAMPLE_RATE,
            n_mfcc=Config.N_MFCC,
            n_chroma=Config.N_CHROMA
        )

        self._initialized = True
        
    def validate_file(self, filepath: str) -> bool:
        if not os.path.exists(filepath):
            raise FileNotFoundError(
                f"File not found: {filepath}\n"
                "Please provide a valid path to a .wav file."
            )

        if not filepath.lower().endswith('.wav'):
            raise ValueError(
                f"Unsupported file format: {filepath}\n"
                "Only .wav files are supported."
            )

        if os.path.getsize(filepath) == 0:
            raise ValueError(f"File is empty: {filepath}")

        return True

    def extract_features(self, filepath: str) -> Optional[np.ndarray]:
        features = self.feature_extractor.process_file(filepath)

        if features is None:
            raise RuntimeError(
                f"Failed to extract features from: {filepath}\n"
                "The audio file may be corrupt or invalid."
            )

        return features

    def preprocess(self, features: np.ndarray) -> np.ndarray:
        features = features.reshape(1, -1)
        features_scaled = self.scaler.transform(features)

        if self.use_pca and self.pca is not None:
            features_processed = self.pca.transform(features_scaled)
        else:
            features_processed = features_scaled

        return features_processed

    def predict(self, filepath: str) -> Tuple[str, float]:
        if not self._initialized:
            self.initialize()

        self.validate_file(filepath)
    

        features = self.extract_features(filepath)
        

        features_processed = self.preprocess(features)
        

        prediction = self.model.predict(features_processed)[0]
        probabilities = self.model.predict_proba(features_processed)[0]

        confidence = probabilities[prediction] * 100
        label = "Parkinson Detected" if prediction == 1 else "Healthy"

        return label, confidence

    def predict_with_details(self, filepath: str) -> dict:
        if not self._initialized:
            self.initialize()

        label, confidence = self.predict(filepath)

        features = self.feature_extractor.process_file(filepath)
        features_processed = self.preprocess(features)
        probabilities = self.model.predict_proba(features_processed)[0]

        actual_model_name = self.model_name
        if not actual_model_name:
            try:
                with open(Config.BEST_MODEL_FILE, 'r') as f:
                    actual_model_name = f.read().strip()
            except:
                actual_model_name = 'SVM'

        return {
            'file': filepath,
            'model': actual_model_name,
            'prediction': label,
            'confidence': confidence,
            'probabilities': {
                'healthy': probabilities[0] * 100,
                'parkinson': probabilities[1] * 100
            },
            'use_pca': self.use_pca
        }


def format_output(result: dict, verbose: bool = False) -> str:
    lines = []

   
    lines.append("PARKINSON'S DISEASE DETECTION - PREDICTION")
    

    lines.append(f"Using Model: {result['model']}")
    lines.append(f"Prediction: {result['prediction']}")
    lines.append(f"Confidence: {result['confidence']:.{Config.CONFIDENCE_DECIMALS}f}%")

    if verbose:
        lines.append("-" * 50)
        lines.append("Detailed Probabilities:")
        lines.append(f"  Healthy:     {result['probabilities']['healthy']:.2f}%")
        lines.append(f"  Parkinson:   {result['probabilities']['parkinson']:.2f}%")
        lines.append(f"PCA Applied:   {result['use_pca']}")

    

    return '\n'.join(lines)


def print_result(result: dict, verbose: bool = False) -> None:
    print(format_output(result, verbose))


def predict(file_path: str, model_name: Optional[str] = None,
            models_dir: str = Config.MODELS_DIR,
            verbose: bool = False) -> dict:
    predictor = ParkinsonPredictor(model_name, models_dir)
    result = predictor.predict_with_details(file_path)
    print_result(result, verbose)
    return result


def main():
    parser = argparse.ArgumentParser(
        description='Predict Parkinson\'s disease from voice audio',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python predict.py --file audio.wav
    python predict.py -f sample.wav --model RandomForest
    python predict.py --file voice.wav --verbose
    python predict.py --record              # Record from mic and predict
    python predict.py --record --duration 5

Models:
    - RandomForest (or RandomForest_PCA)
    - LogisticRegression (or LogisticRegression_PCA)
    - SVM (or SVM_PCA)

If no model is specified, the best model from training is used.
        """
    )

    parser.add_argument(
        '--file', '-f',
        type=str,
        default=None,
        help='Path to the .wav audio file for prediction (not required if using --record)'
    )

    parser.add_argument(
        '--model', '-m',
        type=str,
        default=None,
        help='Model to use (default: best model from best_model.txt)'
    )

    parser.add_argument(
        '--models_dir', '-d',
        type=str,
        default='./models',
        help='Directory containing saved models (default: ./models)'
    )

    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Show detailed prediction information'
    )

    parser.add_argument(
        '--record', '-r',
        action='store_true',
        help='Record from microphone instead of using a file'
    )

    parser.add_argument(
        '--duration',
        type=float,
        default=3.0,
        help='Recording duration in seconds (default: 3)'
    )

    args = parser.parse_args()

    if not args.file and not args.record:
        parser.error("Either --file or --record must be provided")

    if not os.path.exists(args.models_dir):
        logger.error(f"Models directory not found: {args.models_dir}")
        logger.error("Please run train.py and pca_train.py first.")
        sys.exit(1)

    if args.record:
        if not PYAUDIO_AVAILABLE:
            logger.error("\n" + "=" * 60)
            logger.error("PyAudio is not installed!")
            logger.error("=" * 60)
            logger.error("\nTo install PyAudio:")
            logger.error("  macOS:   brew install portaudio && pip install pyaudio")
            logger.error("  Linux:   sudo apt-get install portaudio19-dev && pip install pyaudio")
            logger.error("  Windows: pip install pyaudio")
            logger.error("=" * 60)
            sys.exit(1)

        temp_file = None
        try:
            temp_file = tempfile.NamedTemporaryFile(suffix='.wav', delete=False)
            temp_file.close()

            
            logger.info("VOICE RECORDING")
          

            record_audio(args.duration, temp_file.name)
            args.file = temp_file.name
            logger.info(f"Recording saved to: {temp_file.name}")

        except KeyboardInterrupt:
            logger.info("\nRecording cancelled by user.")
            if temp_file and os.path.exists(temp_file.name):
                os.remove(temp_file.name)
            sys.exit(130)

    try:
        result = predict(
            file_path=args.file,
            model_name=args.model,
            models_dir=args.models_dir,
            verbose=args.verbose
        )

        if args.record and temp_file and os.path.exists(temp_file.name):
            os.remove(temp_file.name)

        if 'Parkinson' in result['prediction']:
            sys.exit(1)
        else:
            sys.exit(0)

    except FileNotFoundError as e:
        logger.error(f"File Error: {e}")
        sys.exit(2)

    except ValueError as e:
        logger.error(f"Validation Error: {e}")
        sys.exit(3)

    except RuntimeError as e:
        logger.error(f"Processing Error: {e}")
        sys.exit(4)

    except Exception as e:
        logger.error(f"Unexpected Error: {str(e)}")
        sys.exit(99)


def record_audio(duration: float, output_path: str, sample_rate: int = 22050) -> None:
    audio = pyaudio.PyAudio()

    stream = audio.open(
        format=pyaudio.paInt16,
        channels=1,
        rate=sample_rate,
        input=True,
        frames_per_buffer=1024
    )

    logger.info(f"Recording for {duration} seconds...")

    logger.info("RECORDING! Speak now...")

    frames = []
    total_chunks = int(sample_rate / 1024 * duration)

    for i in range(total_chunks):
        data = stream.read(1024, exception_on_overflow=False)
        frames.append(data)

        # progress = (i + 1) / total_chunks * 100
        # if progress % 25 < 10:
        #     logger.info(f"  Recording: {progress:.0f}%")

    stream.stop_stream()
    stream.close()
    audio.terminate()

    logger.info("Recording complete!")

    with wave.open(output_path, 'wb') as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(b''.join(frames))


if __name__ == '__main__':
    main()