08 - Deploying YOLO26 on Edge: Raspberry Pi and Jetson

import torch
import torch.quantization as quant
from torch.ao.quantization import get_default_qconfig, prepare, convert
from torchvision import models
import copy

def quantize_model_ptq(
    model: torch.nn.Module,
    calibration_loader,
    backend: str = 'qnnpack'  # 'qnnpack' for ARM, 'x86' for Intel CPU
) -> torch.nn.Module:
    """
    Post-Training Quantization (PTQ): quantize model without retraining.
    Only requires a small calibration dataset (~100-1000 images).

    Flow:
    1. Fuse operations (Conv+BN+ReLU -> single op)
    2. Insert observers for calibration
    3. Run calibration (forward pass on calibration dataset)
    4. Convert to quantized model
    """
    torch.backends.quantized.engine = backend

    model_to_quantize = copy.deepcopy(model)
    model_to_quantize.eval()

    # Step 1: Fuse common layers for efficiency
    model_to_quantize = torch.quantization.fuse_modules(
        model_to_quantize,
        [['conv1', 'bn1', 'relu']],  # adapt to your model's layer names
        inplace=True
    )

    # Step 2: Set qconfig and prepare for calibration
    qconfig = get_default_qconfig(backend)
    model_to_quantize.qconfig = qconfig
    prepared_model = prepare(model_to_quantize, inplace=False)

    # Step 3: Calibration with real data
    print("Running quantization calibration...")
    prepared_model.eval()
    with torch.no_grad():
        for i, (images, _) in enumerate(calibration_loader):
            prepared_model(images)
            if i >= 99:  # 100 calibration batches sufficient
                break
            if i % 10 == 0:
                print(f"  Batch {i+1}/100")

    # Step 4: Convert to quantized model
    quantized_model = convert(prepared_model, inplace=False)

    # Verify size reduction
    def model_size_mb(m: torch.nn.Module) -> float:
        param_size = sum(p.nelement() * p.element_size() for p in m.parameters())
        buffer_size = sum(b.nelement() * b.element_size() for b in m.buffers())
        return (param_size + buffer_size) / (1024 ** 2)

    original_size = model_size_mb(model)
    quantized_size = model_size_mb(quantized_model)
    print(f"Original size: {original_size:.1f} MB")
    print(f"Quantized size: {quantized_size:.1f} MB")
    print(f"Reduction: {original_size / quantized_size:.1f}x")

    return quantized_model


def compare_inference_speed(original_model, quantized_model,
                             input_tensor: torch.Tensor, n_runs: int = 100) -> dict:
    """Compare inference speed between original and quantized model."""
    import time

    results = {}

    for name, model in [('FP32', original_model), ('INT8', quantized_model)]:
        model.eval()
        # Warmup
        with torch.no_grad():
            for _ in range(10):
                model(input_tensor)

        # Benchmark
        start = time.perf_counter()
        with torch.no_grad():
            for _ in range(n_runs):
                model(input_tensor)
        elapsed = time.perf_counter() - start

        avg_ms = (elapsed / n_runs) * 1000
        results[name] = avg_ms
        print(f"{name}: {avg_ms:.2f}ms / inference")

    speedup = results['FP32'] / results['INT8']
    print(f"INT8 Speedup: {speedup:.2f}x")
    return results

from ultralytics import YOLO

model = YOLO('yolo26n.pt')  # nano variant for edge

# ---- TFLite INT8 for Raspberry Pi / Coral TPU ----
model.export(
    format='tflite',
    imgsz=320,        # reduced resolution for edge
    int8=True,        # INT8 quantization
    data='coco.yaml'  # dataset for PTQ calibration
)
# Output: yolo26n_int8.tflite

# ---- NCNN for ARM CPU (Raspberry Pi, Android) ----
model.export(
    format='ncnn',
    imgsz=320,
    half=False  # NCNN uses native FP32 or INT8
)
# Output: yolo26n_ncnn_model/

# ---- TensorRT FP16 for Jetson ----
model.export(
    format='engine',
    imgsz=640,
    half=True,       # FP16
    workspace=2,     # GB workspace (reduced for Jetson Nano)
    device=0
)
# Output: yolo26n.engine

# ---- ONNX + ONNX Runtime for CPU/NPU ----
model.export(
    format='onnx',
    imgsz=320,
    opset=17,
    simplify=True,
    dynamic=False    # fixed batch size for edge deployment
)

print("Export completed for all edge targets")

# ============================================
# RASPBERRY PI 5 SETUP for Computer Vision
# ============================================

# 1. Install base dependencies
# sudo apt update && sudo apt install -y python3-pip libopencv-dev
# pip install ultralytics ncnn onnxruntime

# 2. System optimizations for AI
# In /boot/firmware/config.txt:
# gpu_mem=256           # Increase GPU memory (VideoCore VII)
# over_voltage=6        # Mild overclock
# arm_freq=2800         # Max CPU frequency (stock 2.4GHz)

# ============================================
# INFERENCE with NCNN on Raspberry Pi
# ============================================

import ncnn
import cv2
import numpy as np
import time

class YOLOncnn:
    """
    YOLO inference with NCNN - optimized for ARM CPU.
    NCNN by Tencent is the fastest runtime for ARM CPU.
    """

    def __init__(self, param_path: str, bin_path: str,
                 num_threads: int = 4, input_size: int = 320):
        self.net = ncnn.Net()
        self.net.opt.num_threads = num_threads  # use all cores
        self.net.opt.use_vulkan_compute = False  # no discrete GPU on RPi
        self.net.load_param(param_path)
        self.net.load_model(bin_path)
        self.input_size = input_size

    def predict(self, img_bgr: np.ndarray, conf_thresh: float = 0.4) -> list[dict]:
        """NCNN inference on ARM CPU."""
        h, w = img_bgr.shape[:2]

        # Resize + normalization for NCNN
        img_resized = cv2.resize(img_bgr, (self.input_size, self.input_size))
        img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)

        mat_in = ncnn.Mat.from_pixels(
            img_rgb, ncnn.Mat.PixelType.PIXEL_RGB,
            self.input_size, self.input_size
        )
        mean_vals = [0.485 * 255, 0.456 * 255, 0.406 * 255]
        norm_vals = [1/0.229/255, 1/0.224/255, 1/0.225/255]
        mat_in.substract_mean_normalize(mean_vals, norm_vals)

        ex = self.net.create_extractor()
        ex.input("images", mat_in)
        _, mat_out = ex.extract("output0")

        return self._parse_output(mat_out, conf_thresh, w, h)

    def _parse_output(self, mat_out, conf_thresh,
                      orig_w, orig_h) -> list[dict]:
        """Parse NCNN output into detection format."""
        detections = []
        for i in range(mat_out.h):
            row = np.array(mat_out.row(i))
            confidence = row[4]
            if confidence < conf_thresh:
                continue

            class_scores = row[5:]
            class_id = int(np.argmax(class_scores))
            class_conf = confidence * class_scores[class_id]

            if class_conf >= conf_thresh:
                cx, cy, bw, bh = row[:4]
                x1 = int((cx - bw/2) * orig_w / self.input_size)
                y1 = int((cy - bh/2) * orig_h / self.input_size)
                x2 = int((cx + bw/2) * orig_w / self.input_size)
                y2 = int((cy + bh/2) * orig_h / self.input_size)

                detections.append({
                    'class_id': class_id,
                    'confidence': float(class_conf),
                    'bbox': (x1, y1, x2, y2)
                })

        return detections


def run_rpi_detection_loop(model_param: str, model_bin: str,
                            camera_id: int = 0) -> None:
    """Real-time detection loop optimized for Raspberry Pi."""
    detector = YOLOncnn(model_param, model_bin,
                        num_threads=4, input_size=320)
    cap = cv2.VideoCapture(camera_id)

    # Optimize capture for RPi
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    cap.set(cv2.CAP_PROP_FPS, 30)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # minimize latency

    frame_skip = 2   # Process 1 frame out of 3 to save CPU
    frame_count = 0
    cached_dets = []
    fps_history = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        t0 = time.perf_counter()

        if frame_count % frame_skip == 0:
            cached_dets = detector.predict(frame, conf_thresh=0.4)

        elapsed = time.perf_counter() - t0
        fps = 1.0 / elapsed if elapsed > 0 else 0
        fps_history.append(fps)

        # Visualization
        for det in cached_dets:
            x1, y1, x2, y2 = det['bbox']
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, f"{det['confidence']:.2f}",
                       (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX,
                       0.5, (0,255,0), 2)

        avg_fps = sum(fps_history[-30:]) / min(len(fps_history), 30)
        cv2.putText(frame, f"FPS: {avg_fps:.1f}", (10, 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow('RPi Detection', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        frame_count += 1

    cap.release()
    cv2.destroyAllWindows()
    print(f"Average FPS: {sum(fps_history)/len(fps_history):.1f}")

from ultralytics import YOLO
import cv2
import time

def setup_jetson_pipeline(model_path: str = 'yolo26n.pt') -> YOLO:
    """
    Optimal setup for Jetson Orin:
    1. Export to TensorRT FP16
    2. Configure jetson_clocks for maximum performance
    3. Set performance mode for GPU
    """
    import subprocess

    # Maximize Jetson performance (run once, requires sudo)
    # subprocess.run(['sudo', 'jetson_clocks'], check=True)
    # subprocess.run(['sudo', 'nvpmodel', '-m', '0'], check=True)  # MAXN mode

    model = YOLO(model_path)

    print("Exporting to TensorRT FP16...")
    model.export(
        format='engine',
        imgsz=640,
        half=True,       # FP16 - nearly same accuracy as FP32, 2x faster
        workspace=2,     # GB GPU workspace (Jetson Orin Nano: 8GB shared)
        device=0,
        batch=1,
        simplify=True
    )

    # Load the TensorRT model
    trt_model = YOLO('yolo26n.engine')
    print("TensorRT model ready")
    return trt_model


def run_jetson_pipeline(model: YOLO, source=0) -> None:
    """Real-time pipeline optimized for Jetson with performance stats."""
    cap = cv2.VideoCapture(source)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

    fps_list = []
    frame_count = 0

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break

            t0 = time.perf_counter()
            results = model.predict(
                frame, conf=0.35, iou=0.45,
                verbose=False, half=True  # FP16 inference
            )
            elapsed = time.perf_counter() - t0
            fps = 1.0 / elapsed
            fps_list.append(fps)

            # Annotate with performance info
            annotated = results[0].plot()
            avg_fps = sum(fps_list[-30:]) / min(len(fps_list), 30)

            info_lines = [
                f"FPS: {fps:.0f} (avg: {avg_fps:.0f})",
                f"Detections: {len(results[0].boxes)}",
                f"Inference: {elapsed*1000:.1f}ms"
            ]
            for i, text in enumerate(info_lines):
                cv2.putText(annotated, text, (10, 30 + i * 30),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

            cv2.imshow('Jetson Pipeline', annotated)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

            frame_count += 1

    finally:
        cap.release()
        cv2.destroyAllWindows()
        if fps_list:
            print(f"\n=== Jetson Stats ===")
            print(f"Frames: {frame_count}")
            print(f"Average FPS: {sum(fps_list)/len(fps_list):.1f}")
            print(f"Peak FPS: {max(fps_list):.1f}")
            print(f"Min latency: {1000/max(fps_list):.1f}ms")

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

def apply_structured_pruning(model: nn.Module,
                               amount: float = 0.3,
                               n: int = 2) -> nn.Module:
    """
    Structured L2-norm pruning: removes entire filters/neurons.
    Produces faster inference models (unlike unstructured pruning
    which only produces smaller but not necessarily faster models).

    amount: fraction of filters to remove (0.3 = 30%)
    n: L_n norm used for filter ranking
    """
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            prune.ln_structured(
                module,
                name='weight',
                amount=amount,
                n=n,
                dim=0  # dim=0 = prune output filters
            )
        elif isinstance(module, nn.Linear):
            prune.ln_structured(
                module,
                name='weight',
                amount=amount,
                n=n,
                dim=0
            )

    return model


def remove_pruning_masks(model: nn.Module) -> nn.Module:
    """
    Make pruning permanent: remove masks and 'orig' parameters,
    keeping only the pruned weights. Required before export.
    """
    for name, module in model.named_modules():
        if isinstance(module, (nn.Conv2d, nn.Linear)):
            try:
                prune.remove(module, 'weight')
            except ValueError:
                pass
    return model


def prune_and_finetune(model: nn.Module, train_loader, val_loader,
                        prune_amount: float = 0.2,
                        finetune_epochs: int = 5) -> nn.Module:
    """
    Complete pipeline:
    1. Prune the model (remove prune_amount% of filters)
    2. Fine-tune to recover lost accuracy
    3. Remove masks and finalize
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    print(f"Applying {prune_amount*100:.0f}% structured pruning...")
    model = apply_structured_pruning(model, amount=prune_amount)

    # Brief fine-tuning for accuracy recovery
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(finetune_epochs):
        model.train()
        total_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            loss = criterion(model(images), labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        model.eval()
        correct = total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                preds = model(images).argmax(1)
                correct += preds.eq(labels).sum().item()
                total += labels.size(0)

        print(f"  FT Epoch {epoch+1}/{finetune_epochs} | "
              f"Loss: {total_loss/len(train_loader):.4f} | "
              f"Acc: {100.*correct/total:.2f}%")

    model = remove_pruning_masks(model)
    print("Pruning completed and finalized")
    return model

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """
    Combined loss for knowledge distillation:
    L_total = alpha * L_CE(student, labels) + (1-alpha) * L_KD(student, teacher)

    L_KD = KL divergence between soft predictions (temperature-scaled)
    Temperature T > 1 softens the distributions, revealing inter-class structure.
    """

    def __init__(self, temperature: float = 4.0, alpha: float = 0.3):
        super().__init__()
        self.T = temperature
        self.alpha = alpha
        self.ce = nn.CrossEntropyLoss()

    def forward(self, student_logits: torch.Tensor,
                teacher_logits: torch.Tensor,
                labels: torch.Tensor) -> torch.Tensor:
        # Standard cross-entropy loss
        loss_ce = self.ce(student_logits, labels)

        # Soft prediction loss (KL divergence)
        student_soft = F.log_softmax(student_logits / self.T, dim=1)
        teacher_soft = F.softmax(teacher_logits / self.T, dim=1)
        loss_kd = F.kl_div(student_soft, teacher_soft,
                            reduction='batchmean') * (self.T ** 2)

        return self.alpha * loss_ce + (1 - self.alpha) * loss_kd


def train_with_distillation(teacher_model: nn.Module,
                              student_model: nn.Module,
                              train_loader,
                              epochs: int = 30,
                              temperature: float = 4.0) -> nn.Module:
    """Train student model guided by teacher model."""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    teacher_model.to(device).eval()  # Teacher is frozen
    student_model.to(device)

    criterion = DistillationLoss(temperature=temperature, alpha=0.3)
    optimizer = torch.optim.AdamW(student_model.parameters(),
                                   lr=1e-3, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=epochs
    )

    for epoch in range(epochs):
        student_model.train()
        total_loss = 0.0

        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)

            with torch.no_grad():
                teacher_logits = teacher_model(images)

            student_logits = student_model(images)
            loss = criterion(student_logits, teacher_logits, labels)

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        scheduler.step()
        print(f"Epoch {epoch+1}/{epochs} | "
              f"Loss: {total_loss/len(train_loader):.4f} | "
              f"LR: {scheduler.get_last_lr()[0]:.2e}")

    return student_model

import onnxruntime as ort
import numpy as np
import cv2
import time
from typing import Optional

class ONNXInferenceEngine:
    """
    Cross-platform ONNX Runtime inference.
    Automatically selects best execution provider:
    TensorRT > CUDA > CPU
    """

    EXECUTION_PROVIDERS = [
        'TensorrtExecutionProvider',   # Jetson with TensorRT
        'CUDAExecutionProvider',        # NVIDIA GPU (generic)
        'CPUExecutionProvider'          # Fallback: any device
    ]

    def __init__(self, model_path: str,
                 providers: Optional[list] = None):
        providers = providers or self.EXECUTION_PROVIDERS

        # Session options for optimization
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = (
            ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        )
        sess_options.intra_op_num_threads = 4
        sess_options.inter_op_num_threads = 1

        # Check available providers
        available = ort.get_available_providers()
        selected = [p for p in providers if p in available]
        print(f"Available providers: {available}")
        print(f"Using: {selected[0]}")

        self.session = ort.InferenceSession(
            model_path,
            sess_options=sess_options,
            providers=selected
        )

        # Get I/O shapes
        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        print(f"Input: {self.input_name} {self.input_shape}")

    def preprocess(self, img_bgr: np.ndarray,
                   input_size: int = 320) -> np.ndarray:
        """Preprocess image for ONNX model."""
        img = cv2.resize(img_bgr, (input_size, input_size))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        img = np.expand_dims(img, 0)        # CHW -> NCHW
        return np.ascontiguousarray(img)

    def infer(self, img_bgr: np.ndarray,
              input_size: int = 320) -> list:
        """Run inference and return raw outputs."""
        input_data = self.preprocess(img_bgr, input_size)
        outputs = self.session.run(
            None,
            {self.input_name: input_data}
        )
        return outputs

    def benchmark(self, input_size: int = 320,
                  n_runs: int = 100) -> dict:
        """Benchmark inference speed."""
        dummy_img = np.random.randint(0, 255,
                                       (480, 640, 3),
                                       dtype=np.uint8)
        # Warmup
        for _ in range(10):
            self.infer(dummy_img, input_size)

        times = []
        for _ in range(n_runs):
            t0 = time.perf_counter()
            self.infer(dummy_img, input_size)
            times.append(time.perf_counter() - t0)

        avg_ms = np.mean(times) * 1000
        p99_ms = np.percentile(times, 99) * 1000
        print(f"Avg latency: {avg_ms:.2f}ms")
        print(f"P99 latency: {p99_ms:.2f}ms")
        print(f"Avg FPS: {1000/avg_ms:.1f}")
        return {"avg_ms": avg_ms, "p99_ms": p99_ms}