seavision-python 0.1.1__py3-none-any.whl

seavision/engine/detectors/yolo/detector.py

@@ -0,0 +1,236 @@
+"""YOLO detector implementation using Ultralytics for SeaVision.
+
+This module wraps Ultralytics YOLO to implement the DetectorBase interface, so
+it can be used in the SeaVision detection pipeline as a pure, stateless
+detector (no built-in tracking or persistence filtering).
+
+Example:
+    from seavision.engine.detectors.yolo import YOLODetector, YOLODetectorConfig
+
+    config = YOLODetectorConfig(
+        model_path="./models/yolo11n.pt",
+        conf_threshold=0.4,
+        iou_threshold=0.5,
+    )
+
+    with YOLODetector(config) as detector:
+        for frame, context in source.iter_frames():
+            for det in detector.process_frame(frame, context):
+                conf = det.confidence if det.confidence is not None else float("nan")
+                print(
+                    f"{det.source_file} "
+                    f"frame={det.frame_number} "
+                    f"label={det.label} "
+                    f"conf={conf:.2f}"
+                )
+"""
+
+import logging
+import os
+from typing import Iterator, Optional
+
+import numpy as np
+import ultralytics
+
+from ...source import FrameContext
+from ..base import Detection, DetectorBase
+from .config import YOLODetectorConfig
+
+logger = logging.getLogger(__name__)
+
+
+class YOLODetector(DetectorBase):
+    """Object detector based on Ultralytics YOLO.
+
+    This detector implements the DetectorBase interface and produces
+    :class:`Detection` objects directly from YOLO predictions. It does **not**
+    perform any temporal tracking or persistence filtering; those can be
+    composed separately if needed.
+    """
+
+    def __init__(self, config: Optional[YOLODetectorConfig] = None):
+        """Initialise the YOLO detector.
+
+        Args:
+            config: Optional :class:`YOLODetectorConfig`. If None, defaults are
+                used.
+        """
+
+        # This import will already have been attempted at module import time,
+        # but we keep the reference here for clarity and to make the
+        # dependency explicit.
+        if not hasattr(ultralytics, "YOLO"):
+            raise ImportError(
+                "Ultralytics YOLO is required for YOLODetector. "
+                "Install with: pip install ultralytics>=8.3.237"
+            )
+
+        self.config = config or YOLODetectorConfig()
+
+        # Ultralytics model instance (lazy-loaded)
+        self._model: Optional[ultralytics.YOLO] = None
+
+        logger.info(
+            "YOLODetector initialised (model=%s, device=%s)",
+            self.config.model_path,
+            self.config.device,
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _ensure_model_loaded(self) -> None:
+        """Lazily load the YOLO model from a local path.
+
+        This method is idempotent. It also enforces that the model path exists
+        locally to avoid implicit downloads or side-effectful behaviour.
+        """
+
+        if self._model is not None:
+            return
+
+        if not os.path.exists(self.config.model_path):
+            raise FileNotFoundError(
+                "YOLO weights not found at: "
+                f"{self.config.model_path}. "
+                "Download the weights separately and set model_path to the "
+                "local file."
+            )
+
+        logger.info("Loading YOLO model from %s", self.config.model_path)
+        self._model = ultralytics.YOLO(self.config.model_path)
+
+        # Move to device if specified
+        if self.config.device:
+            try:
+                self._model.to(self.config.device)
+            except Exception as exc:
+                logger.warning(
+                    "Failed to move YOLO model to device '%s': %s. "
+                    "Using default device instead.",
+                    self.config.device,
+                    exc,
+                )
+
+        logger.info("YOLO model loaded successfully")
+
+    def _class_name_from_id(self, cls_id: int) -> str:
+        """Map a YOLO class index to a human-readable name, if available."""
+
+        if self._model is None:
+            return str(cls_id)
+
+        names = getattr(self._model, "names", None)
+
+        if isinstance(names, dict):
+            return str(names.get(cls_id, cls_id))
+        if isinstance(names, (list, tuple)):
+            if 0 <= cls_id < len(names):
+                return str(names[cls_id])
+
+        return str(cls_id)
+
+    # ------------------------------------------------------------------
+    # DetectorBase interface
+    # ------------------------------------------------------------------
+
+    def process_frame(
+        self,
+        frame: np.ndarray,
+        context: FrameContext,
+    ) -> Iterator[Detection]:
+        """Process a single video frame and yield YOLO detections.
+
+        Args:
+            frame: BGR image as a numpy array with shape (H, W, 3).
+            context: FrameContext containing metadata for this frame.
+
+        Yields:
+            Detection objects corresponding to YOLO predictions that pass the
+            configured confidence and class filters.
+        """
+
+        self._ensure_model_loaded()
+
+        # Run YOLO inference on the in-memory frame. All saving-related
+        # options are explicitly disabled to avoid side effects.
+        results_list = self._model.predict(
+            frame,
+            imgsz=self.config.imgsz,
+            conf=self.config.conf_threshold,
+            iou=self.config.iou_threshold,
+            max_det=self.config.max_detections,
+            classes=self.config.classes,
+            device=self.config.device,
+            half=self.config.half,
+            verbose=self.config.verbose,
+            save=False,
+            save_frames=False,
+            save_txt=False,
+            save_conf=False,
+            save_crop=False,
+            project=None,
+            name=None,
+        )
+
+        if not results_list:
+            return
+
+        results = results_list[0]
+        boxes = results.boxes
+
+        if boxes is None or len(boxes) == 0:
+            return
+
+        for i in range(len(boxes)):
+            box = boxes[i]
+
+            # xyxy format as numpy array
+            xyxy = box.xyxy[0].cpu().numpy()
+            x1, y1, x2, y2 = xyxy
+
+            width = float(x2 - x1)
+            height = float(y2 - y1)
+            xc = float(x1 + width / 2.0)
+            yc = float(y1 + height / 2.0)
+
+            # Confidence (YOLO already applies conf threshold, but we keep
+            # this for completeness and potential future overrides).
+            confidence: Optional[float] = None
+            if box.conf is not None and len(box.conf) > 0:
+                confidence = float(box.conf[0])
+                if confidence < self.config.conf_threshold:
+                    continue
+
+            # Class index
+            cls_id: Optional[int] = None
+            if box.cls is not None and len(box.cls) > 0:
+                cls_id = int(box.cls[0])
+
+            label: Optional[str] = None
+            if self.config.output_labels and cls_id is not None:
+                label = self._class_name_from_id(cls_id)
+
+            yield Detection(
+                source_file=context.source_file,
+                timestamp=context.timestamp,
+                frame_number=context.frame_number,
+                xc=xc,
+                yc=yc,
+                width=width,
+                height=height,
+                confidence=confidence,
+                label=label,
+                track_id=None,
+                mask=None,
+            )
+
+    def reset(self) -> None:
+        """Reset detector state between videos.
+
+        The YOLO model itself remains loaded; this method exists to satisfy
+        the DetectorBase contract and for future extension if needed.
+        """
+
+        logger.debug("YOLODetector reset called (no persistent state to clear)")

seavision/engine/export/__init__.py

@@ -0,0 +1,24 @@
+"""
+Model export for edge deployment.
+
+This package converts trained YOLO .pt weights to optimised formats
+(ONNX, TFLite, NCNN) for deployment on edge devices like Raspberry Pi.
+
+Public API:
+- ExportConfig: Configuration for the export process.
+- ExportTarget: Enum of supported export formats.
+- ExportResult: Result of an export operation.
+- ValidationResult: Result of export validation.
+- ModelExporter: The main export class.
+"""
+
+from .config import ExportConfig, ExportResult, ExportTarget, ValidationResult
+from .exporter import ModelExporter
+
+__all__ = [
+    "ExportConfig",
+    "ExportResult",
+    "ExportTarget",
+    "ModelExporter",
+    "ValidationResult",
+]

seavision/engine/export/config.py

@@ -0,0 +1,160 @@
+"""Configuration and result types for model exprt to edge formats."""
+
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from pathlib import Path
+from typing import Dict, List, Optional
+
+
+class ExportTarget(Enum):
+    """
+    Supported export target formats.
+    
+    Each target produces a different model artifact optimised for a specific
+    runtime or hardware accelerator.
+    """
+
+    ONNX = auto()        # ONNX Runtime - best for Raspberry Pi CPU
+    TFLITE = auto()      # TensorFlow Lite - required for Coral Edge TPU
+    NCNN = auto()        # Tencent ncnn - lightweight, no runtime dependency
+    TORCHSCRIPT = auto() # Torchhscript - fallback requires PyTorch on device
+
+
+# Maps Export Target to the Ultralytics export format string and the file
+# extension (or directory suffix) of the resulting artifact.
+TARGET_EXPORT_INFO: Dict[ExportTarget, Dict[str, str]] = {
+    ExportTarget.ONNX: {
+        "format": "onnx",
+        "suffix": ".onnx",
+        "runtime_package": "onnxruntime",
+    },
+    ExportTarget.TFLITE: {
+        "format": "tflite",
+        "suffix": ".tflite",
+        "runtime_package": "tflite-runtime",
+    },
+    ExportTarget.NCNN: {
+        "format": "ncnn",
+        "suffix": "_ncnn_model",
+        "runtime_package": "ncnn",
+    },
+    ExportTarget.TORCHSCRIPT: {
+        "format": "torchscript",
+        "suffix": ".torchscript",
+        "runtime_package": "torch",
+    },
+}
+
+
+@dataclass
+class ExportConfig:
+    """
+    Configuration for exporting a trained YOLO model to an edge format.
+    
+    Attributes:
+        weights_path: Path to the trained .pt weights file. Must exist.
+        output_dir: Directory to write the exported model and metadata. Created
+            if it does not exist.
+        target: Export target format. Determines which runtime is needed on the
+            edge device and which optimisations are applied. Default is ONNX -
+            the most reliable and best-performaing options for Raspberry Pi
+            4/5 CPU inference.
+        imgsz: Input image size for the exported model. The edge runtime must
+            use the same size. Smaller sizes (320) run faster but detect fewer
+            small objects. Larger sizes (640) are the more accurate but slower.
+            This is hte single most impactful performance knob for edge 
+            deployment.
+        half: Use float16 quanitisation. Reduces model size by ~50% with minimal
+            accuracy loss on most detection tasks. Not supported on all targets
+            or devices - ONNX on ARM CPUs does not benefit from half precision,
+            so this defaults to False.
+        int8: Use int8 quantisation. Aggressive size and speed optimisation. May
+            impact detection accuracy, especially for small or low-contrast
+            objects. Requires a calibration dataset for TFLite exports.
+        simplify: Run ONNX graph simplification passes. Folds constants,
+            eliminates dead nodes, and merges operations. Reccommended for all
+            edge deployments - the simplified graph runs faster and is smaller.
+            Only applies to ONNX target.
+        opset: ONNX opset version. Higher opsets support more operations but may
+            not be supported by older ONNX Runtime versions. Default 12 is
+            widely compatible. Only applies to ONNX target.
+        validate: Run export validation after conversion. Loads both the
+            original PyTorch model and the exported model, runs a synthetic
+            frame through each, and ocmpares outputs. Catches silent export
+            failures where the conversion completes without error but produces
+            different detections.
+        device: Device to use for the PyTorch model during export and
+            validation. Use 'cpu' unless you have a CUDA GPU available, in which
+            case 'cuda' will speed up the export process.
+    """
+
+    weights_path: str = ""
+    output_dir: str = "./export"
+
+    target: ExportTarget = ExportTarget.ONNX
+    imgsz: int = 640
+    half: bool = False
+    int8: bool = False
+    simplify: bool = True
+    opset: int = 12
+
+    validate: bool = True
+    device: str = "cpu"
+
+
+@dataclass
+class ValidationResult:
+    """Result of comparing PyTorch vs exported model outputs.
+
+    Attributes:
+        passed: Whether the validation passed — the exported model
+            produces detections consistent with the PyTorch original.
+        num_detections_pytorch: Number of detections from the PyTorch model.
+        num_detections_exported: Number of detections from the exported model.
+        max_box_deviation: Maximum pixel deviation between matched boxes.
+            A value under 2.0 pixels is normal and caused by floating
+            point differences between runtimes.
+        max_confidence_deviation: Maximum confidence score deviation
+            between matched detections.
+        details: Human-readable summary of the comparison.
+    """
+
+    passed: bool = False
+    num_detections_pytorch: int = 0
+    num_detections_exported: int = 0
+    max_box_deviation: float = 0.0
+    max_confidence_deviation: float = 0.0
+    details: str = ""
+
+
+@dataclass
+class ExportResult:
+    """Result of a model export operation.
+
+    Returned by ModelExporter.export(). Contains everything the bundle
+    builder needs to package the model for deployment.
+
+    Attributes:
+        success: Whether the export completed without error.
+        exported_path: Path to the exported model file or directory.
+        config: The ExportConfig that was used.
+        model_size_mb: Size of the exported model in megabytes.
+        num_classes: Number of detection classes in the model.
+        class_names: Mapping of class index to human-readable name,
+            extracted from the YOLO model's metadata.
+        input_shape: Expected input tensor shape, e.g. [1, 3, 640, 640].
+        output_shape: Output tensor shape, e.g. [1, 84, 8400].
+        validation: Validation result, if validation was enabled.
+        error: Error message if the export failed.
+    """
+
+    success: bool = False
+    exported_path: str = ""
+    config: Optional[ExportConfig] = None
+    model_size_mb: float = 0.0
+    num_classes: int = 0
+    class_names: Dict[int, str] = field(default_factory=dict)
+    input_shape: List[int] = field(default_factory=list)
+    output_shape: List[int] = field(default_factory=list)
+    validation: Optional[ValidationResult] = None
+    error: str = ""

seavision/engine/export/exporter.py

@@ -0,0 +1,286 @@
+"""
+Model Exporter - converts trained YOLO .pt weights to edge format.
+
+This module wraps Ultralytics' export pipeline with SeaVision-specific defaults
+and adds export validation. The user interacts with this class; they never need
+to call ultralytics.YOLO.export() directly.
+
+Example:
+    from seavision.engine.export import ModelExporter, ExportConfig
+
+    config = ExportConfig(
+        weights_path="models/fidh_detector.pt",
+        output_dir="./deployment",
+        imgsz=320,
+    )
+
+    exporter = ModelExporter(config)
+    result = exporter.export()
+
+    if result.success:
+        print(f"Exported to: {result.exported_path}")
+        print(f"Model size: {result.model_size_mb:.1f} MB")
+        print(f"Classes: {result.class_names}")
+    else:
+        print(f"Export failed: {result.error}")
+"""
+
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional
+import json
+
+from .config import (
+    ExportConfig,
+    ExportResult,
+    ExportTarget,
+    TARGET_EXPORT_INFO,
+)
+from.validation import validate_export
+
+logger = logging.getLogger(__name__)
+
+
+class ExportError(Exception):
+    """Raised when a model export fails."""
+    pass
+
+
+class ModelExporter:
+    """
+    Converts trained YOLO .pt weights to edge-optimised formats.
+
+    This class is the dingle entrypoint for model export, It:
+
+    1. Loads the PyTorch model via Ultralytics.
+    2. Extracts model metadata (class names, input/output shapes).
+    3. Calls Ultralytics export with SeaVision-specific settings.
+    4. Optionally validates the export by comparing inference outputs.
+    5. Copies the exported artifact to the configured output directory.
+
+    The exporter requires Ultralytics and PyTorch - it runs land-side, not on
+    the edge device.
+    """
+
+    def __init__(self, config: ExportConfig):
+        """
+        Initialise the exporter.
+
+        Args:
+            config: Export configuration
+        """
+        self.config = config
+
+    def export(self) -> ExportResult:
+        """
+        Export the model and return metadata about the result.
+
+        This is the main entrypoint. it handles the full export lifecycle:
+        validation of inputs, model loading, export, optional validation, and
+        artifact copying.
+
+        Returns:
+            ExportResult with success status, exported path, model metadata,
+            and validation results.
+        """
+        config = self.config
+
+        # --- Input validation ---
+        weights = Path(config.weights_path)
+        if not weights.exists():
+            return ExportResult(
+                success=False,
+                config=config,
+                error=f"Weights file not found: {config.weights_path}",
+            )
+        
+        if config.target not in TARGET_EXPORT_INFO:
+            return ExportResult(
+                success=False,
+                config=config,
+                error=f"Unsupported export target: {config.target}",
+            )
+        
+        target_info = TARGET_EXPORT_INFO[config.target]
+        output_dir = Path(config.output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # --- Load model ---
+        try:
+            import ultralytics
+            logger.info("Loading YOLO model from %s", config.weights_path)
+            model = ultralytics.YOLO(config.weights_path)
+        except ImportError:
+            return ExportResult(
+                success=False,
+                config=config,
+                error=(
+                    "Ultralytics is required for model export. "
+                    "Install with: pip install ultralytics"
+                ),
+            )
+        except Exception as exc:
+            return ExportResult(
+                success=False,
+                config=config,
+                error=f"Failed to load model: {exc}",
+            )
+
+        # --- Extract metadata ---
+        class_names: Dict[int, str] = {}
+        names = getattr(model, "names", None)
+        if isinstance(names, dict):
+            class_names = {int(k): str(v) for k, v in names.items()}
+        elif isinstance(names, (list, tuple)):
+            class_names = {i: str(n) for i, n in enumerate(names)}
+        
+        num_classes = len(class_names) if class_names else 0
+
+        # --- Build export kwargs ---
+        export_kwargs = {
+            "format": target_info["format"],
+            "imgsz": config.imgsz,
+            "half": config.half,
+            "int8": config.int8,
+            "device": config.device,
+        }
+
+        if config.target == ExportTarget.ONNX:
+            export_kwargs["simplify"] = config.simplify
+            export_kwargs["opset"] = config.opset
+
+        # --- Run export ---
+        try:
+            logger.info(
+                "Exporting model to %s (imgsz=%d, half=%s, int8=%s)",
+                config.target.name,
+                config.imgsz,
+                config.half,
+                config.int8,
+            )
+            export_path_raw = model.export(**export_kwargs)
+
+            # Ultralytics returns the path as a string
+            export_path = str(export_path_raw)
+            logger.info("Ultralytics export completed: %s", export_path)
+        except Exception as exc:
+            return ExportResult(
+                success=False,
+                config=config,
+                error=f"Export failed: {exc}",
+            )
+
+        # --- Copy to output directory ---
+        src = Path(export_path)
+        suffix = target_info["suffix"]
+        dest_name = weights.stem + suffix
+        dest = output_dir / dest_name
+
+        try:
+            if src.is_dir():
+                if dest.exists():
+                    shutil.rmtree(dest)
+                shutil.copytree(src, dest)
+            else:
+                shutil.copy2(src, dest)
+
+            logger.info("Exported model copied to %s", dest)
+        except Exception as exc:
+            return ExportResult(
+                success=False,
+                config=config,
+                error=f"Failed to copy exported model to output dir: {exc}",
+            )
+        
+        # --- Compute model size ---
+        if dest.is_dir():
+            model_size = sum(f.stat().st_size for f in dest.rglob("*"))
+        else:
+            model_size = dest.stat().st_size
+        model_size_mb = model_size / (1024 * 1024)
+
+        # --- Determine input/output shapes ---
+        # These come from the ONNX metadata if available, otherwise inferred
+        input_shape = [1, 3, config.imgsz, config.imgsz]
+        output_shape: List[int] = []
+
+        if config.target == ExportTarget.ONNX:
+            try:
+                import onnxruntime as ort
+
+                sess = ort.InferenceSession(
+                    str(dest), providers=["CPUExecutionProvider"]
+                )
+                input_shape = list(sess.get_inputs()[0].shape)
+                output_shape = list(sess.get_outputs()[0].shape)
+            except Exception:
+                # Not critical — shapes are informational
+                pass
+
+        # --- Validate ---
+        validation_result = None
+        if config.validate:
+            validation_result = validate_export(config, str(dest))
+
+            if not validation_result.passed:
+                logger.warning(
+                    "Export validation failed. The exported model may "
+                    "produce different results than the original. "
+                    "Details: %s",
+                    validation_result.details,
+                )
+
+        # --- Write metadata ---
+        # --- Write metadata ---
+        self._write_export_metadata(dest, config, class_names, num_classes)
+
+        return ExportResult(
+            success=True,
+            exported_path=str(dest),
+            config=config,
+            model_size_mb=round(model_size_mb, 2),
+            num_classes=num_classes,
+            class_names=class_names,
+            input_shape=input_shape,
+            output_shape=output_shape,
+            validation=validation_result,
+        )
+    
+    def _write_export_metadata(
+        self,
+        dest: Path,
+        config: ExportConfig,
+        class_names: Dict[int, str],
+        num_classes: int,
+    ) -> None:
+        """
+        Write a JSON metadata file alongside the exported model.
+
+        The metadata file contains everything the edge runtime needs to 
+        configure itself: image size, number of classes, class names, and the
+        export target. The bundle builder reads this file; the edge runtime
+        reads it at startup.
+        """
+        metadata = {
+            "seavision_export_version": 1,
+            "source_weights": config.weights_path,
+            "target": config.target.name,
+            "imgsz": config.imgsz,
+            "half": config.half,
+            "int8": config.int8,
+            "num_classes": num_classes,
+            "class_names": class_names,
+        }
+
+        # Place metadata next to the model
+        if dest.is_dir():
+            meta_path = dest / "export_metadata.json"
+        else:
+            meta_path = dest.with_suffix(".json")
+
+        with open(meta_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+
+        logger.info("Export metadata written to %s", meta_path)