inference.py

"""
Inference Script for Anomaly Detection.

Supports PyTorch, ONNX Runtime, and TensorRT backends.
"""

import argparse
import json
import time
from pathlib import Path
from typing import Dict, Union

import numpy as np
import torch
from loguru import logger

# Import inference engines
try:
    import onnxruntime as ort
    ONNX_AVAILABLE = True
except ImportError:
    ONNX_AVAILABLE = False
    logger.warning("ONNX Runtime not available")

try:
    import tensorrt as trt
    import pycuda.driver as cuda
    import pycuda.autoinit
    TENSORRT_AVAILABLE = True
except ImportError:
    TENSORRT_AVAILABLE = False
    logger.warning("TensorRT not available")


class PyTorchInference:
    """PyTorch inference engine."""
    
    def __init__(self, model_path: str, device: str = "cuda"):
        """Initialize PyTorch inference."""
        self.device = device
        
        # Load model
        checkpoint = torch.load(model_path, map_location=device)
        self.model = checkpoint["model"]  # Or load architecture and weights separately
        self.model.eval()
        self.model.to(device)
        
        logger.info(f"PyTorch model loaded from {model_path}")
    
    def predict(self, input_data: np.ndarray) -> Dict[str, Union[float, np.ndarray]]:
        """Run inference."""
        # Convert to tensor
        x = torch.FloatTensor(input_data).to(self.device)
        
        # Add batch dimension if needed
        if x.dim() == 1:
            x = x.unsqueeze(0)
        
        # Inference
        with torch.no_grad():
            start_time = time.time()
            reconstruction = self.model(x)
            inference_time = (time.time() - start_time) * 1000  # ms
        
        # Calculate reconstruction error
        error = torch.mean((x - reconstruction) ** 2, dim=1)
        
        return {
            "reconstruction": reconstruction.cpu().numpy(),
            "reconstruction_error": error.cpu().numpy(),
            "inference_time_ms": inference_time,
        }


class ONNXInference:
    """ONNX Runtime inference engine."""
    
    def __init__(self, model_path: str, device: str = "cuda"):
        """Initialize ONNX inference."""
        if not ONNX_AVAILABLE:
            raise RuntimeError("ONNX Runtime not available")
        
        # Setup providers
        providers = []
        if device == "cuda" and "CUDAExecutionProvider" in ort.get_available_providers():
            providers.append("CUDAExecutionProvider")
        providers.append("CPUExecutionProvider")
        
        # Create session
        self.session = ort.InferenceSession(model_path, providers=providers)
        
        # Get input/output names
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        
        logger.info(f"ONNX model loaded from {model_path}")
        logger.info(f"Using providers: {self.session.get_providers()}")
    
    def predict(self, input_data: np.ndarray) -> Dict[str, Union[float, np.ndarray]]:
        """Run inference."""
        # Add batch dimension if needed
        if input_data.ndim == 1:
            input_data = input_data.reshape(1, -1)
        
        # Convert to float32
        input_data = input_data.astype(np.float32)
        
        # Inference
        start_time = time.time()
        reconstruction = self.session.run(
            [self.output_name],
            {self.input_name: input_data},
        )[0]
        inference_time = (time.time() - start_time) * 1000  # ms
        
        # Calculate reconstruction error
        error = np.mean((input_data - reconstruction) ** 2, axis=1)
        
        return {
            "reconstruction": reconstruction,
            "reconstruction_error": error,
            "inference_time_ms": inference_time,
        }


class TensorRTInference:
    """TensorRT inference engine."""
    
    def __init__(self, model_path: str):
        """Initialize TensorRT inference."""
        if not TENSORRT_AVAILABLE:
            raise RuntimeError("TensorRT not available")
        
        # Load engine
        with open(model_path, "rb") as f:
            engine_data = f.read()
        
        # Create runtime
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        self.engine = self.runtime.deserialize_cuda_engine(engine_data)
        self.context = self.engine.create_execution_context()
        
        # Allocate buffers
        self.inputs = []
        self.outputs = []
        self.bindings = []
        
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding))
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            
            # Allocate host and device buffers
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            
            self.bindings.append(int(device_mem))
            
            if self.engine.binding_is_input(binding):
                self.inputs.append({"host": host_mem, "device": device_mem})
            else:
                self.outputs.append({"host": host_mem, "device": device_mem})
        
        # Create stream
        self.stream = cuda.Stream()
        
        logger.info(f"TensorRT engine loaded from {model_path}")
    
    def predict(self, input_data: np.ndarray) -> Dict[str, Union[float, np.ndarray]]:
        """Run inference."""
        # Add batch dimension if needed
        if input_data.ndim == 1:
            input_data = input_data.reshape(1, -1)
        
        # Copy input to host buffer
        np.copyto(self.inputs[0]["host"], input_data.ravel())
        
        # Transfer input to device
        cuda.memcpy_htod_async(
            self.inputs[0]["device"],
            self.inputs[0]["host"],
            self.stream,
        )
        
        # Run inference
        start_time = time.time()
        self.context.execute_async_v2(
            bindings=self.bindings,
            stream_handle=self.stream.handle,
        )
        
        # Transfer output from device
        cuda.memcpy_dtoh_async(
            self.outputs[0]["host"],
            self.outputs[0]["device"],
            self.stream,
        )
        
        # Synchronize
        self.stream.synchronize()
        inference_time = (time.time() - start_time) * 1000  # ms
        
        # Reshape output
        reconstruction = self.outputs[0]["host"].reshape(input_data.shape)
        
        # Calculate reconstruction error
        error = np.mean((input_data - reconstruction) ** 2, axis=1)
        
        return {
            "reconstruction": reconstruction,
            "reconstruction_error": error,
            "inference_time_ms": inference_time,
        }


def create_inference_engine(
    model_path: str,
    backend: str = "pytorch",
    device: str = "cuda",
):
    """
    Create inference engine based on backend.
    
    Args:
        model_path: Path to model file
        backend: Backend type ('pytorch', 'onnx', 'tensorrt')
        device: Device to use ('cuda' or 'cpu')
        
    Returns:
        Inference engine
    """
    backend = backend.lower()
    
    if backend == "pytorch":
        return PyTorchInference(model_path, device)
    elif backend == "onnx":
        return ONNXInference(model_path, device)
    elif backend == "tensorrt":
        return TensorRTInference(model_path)
    else:
        raise ValueError(f"Unknown backend: {backend}")


def detect_anomaly(
    reconstruction_error: float,
    threshold: float,
    method: str = "threshold",
) -> bool:
    """
    Detect if sample is anomalous.
    
    Args:
        reconstruction_error: Reconstruction error value
        threshold: Anomaly threshold
        method: Detection method
        
    Returns:
        True if anomalous, False otherwise
    """
    if method == "threshold":
        return reconstruction_error > threshold
    else:
        return False


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Run anomaly detection inference")
    
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Path to model file",
    )
    
    parser.add_argument(
        "--backend",
        type=str,
        default="onnx",
        choices=["pytorch", "onnx", "tensorrt"],
        help="Inference backend",
    )
    
    parser.add_argument(
        "--input",
        type=str,
        help="Path to input data file (.npy)",
    )
    
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.1,
        help="Anomaly detection threshold",
    )
    
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device to use",
    )
    
    parser.add_argument(
        "--realtime",
        action="store_true",
        help="Real-time monitoring mode",
    )
    
    parser.add_argument(
        "--sensor",
        type=str,
        help="Sensor device path for real-time mode",
    )
    
    parser.add_argument(
        "--output",
        type=str,
        help="Output file for results (.json)",
    )
    
    return parser.parse_args()


def main():
    """Main inference function."""
    args = parse_args()
    
    logger.info(f"Loading model: {args.model}")
    logger.info(f"Backend: {args.backend}")
    
    # Create inference engine
    engine = create_inference_engine(
        model_path=args.model,
        backend=args.backend,
        device=args.device,
    )
    
    # Real-time mode
    if args.realtime:
        logger.info("Real-time monitoring mode")
        if not args.sensor:
            logger.error("Sensor device path required for real-time mode")
            return
        
        logger.info(f"Monitoring sensor: {args.sensor}")
        # TODO: Implement real-time sensor reading
        logger.warning("Real-time mode not yet implemented")
        return
    
    # Batch inference mode
    if not args.input:
        logger.error("Input file required for batch mode")
        return
    
    logger.info(f"Loading input data from {args.input}")
    input_data = np.load(args.input)
    
    logger.info(f"Input shape: {input_data.shape}")
    
    # Handle multiple samples
    if input_data.ndim == 1:
        input_data = input_data.reshape(1, -1)
    
    # Run inference
    results = []
    total_time = 0
    
    for i, sample in enumerate(input_data):
        result = engine.predict(sample)
        
        error = result["reconstruction_error"][0]
        is_anomaly = detect_anomaly(error, args.threshold)
        
        total_time += result["inference_time_ms"]
        
        sample_result = {
            "sample_id": i,
            "reconstruction_error": float(error),
            "is_anomaly": bool(is_anomaly),
            "inference_time_ms": result["inference_time_ms"],
        }
        
        results.append(sample_result)
        
        logger.info(
            f"Sample {i}: Error={error:.6f}, "
            f"Anomaly={is_anomaly}, Time={result['inference_time_ms']:.2f}ms"
        )
    
    # Summary
    num_samples = len(results)
    num_anomalies = sum(r["is_anomaly"] for r in results)
    avg_time = total_time / num_samples
    
    summary = {
        "num_samples": num_samples,
        "num_anomalies": num_anomalies,
        "anomaly_rate": num_anomalies / num_samples,
        "avg_inference_time_ms": avg_time,
        "throughput_samples_per_sec": 1000.0 / avg_time if avg_time > 0 else 0,
    }
    
    logger.info("\n=== Summary ===")
    logger.info(f"Samples processed: {summary['num_samples']}")
    logger.info(f"Anomalies detected: {summary['num_anomalies']} ({summary['anomaly_rate']:.2%})")
    logger.info(f"Avg inference time: {summary['avg_inference_time_ms']:.2f} ms")
    logger.info(f"Throughput: {summary['throughput_samples_per_sec']:.1f} samples/s")
    
    # Save results
    if args.output:
        output_data = {
            "summary": summary,
            "results": results,
        }
        
        with open(args.output, "w") as f:
            json.dump(output_data, f, indent=2)
        
        logger.info(f"Results saved to {args.output}")


if __name__ == "__main__":
    main()