oriented-det 0.1.0__py3-none-any.whl

export/tests/test_faster_rcnn_export_parity.py

@@ -0,0 +1,201 @@
+"""Parity: PyTorch full detect vs pre-NMS export + postprocess."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_EXPORT_DIR = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+if str(_EXPORT_DIR) not in sys.path:
+    sys.path.insert(0, str(_EXPORT_DIR))
+
+from export.postprocess import finalize_detections_numpy, meta_to_finalize_kwargs  # noqa: E402
+import wrappers as _wrappers  # noqa: E402
+from oriented_det import RotatedFasterRCNN  # noqa: E402
+from oriented_det.models.faster_rcnn_inference import faster_rcnn_inference  # noqa: E402
+from oriented_det.models.utils import rboxes_to_tensor  # noqa: E402
+
+RotatedFasterRCNNPreNmsExportWrapper = _wrappers.RotatedFasterRCNNPreNmsExportWrapper
+
+_DEPLOY_CONFIG = _REPO_ROOT / "deploy/app/config.json"
+_DEPLOY_CKPT = _REPO_ROOT / "deploy/app/weights/model.pth"
+
+
+@pytest.fixture
+def tiny_faster_rcnn() -> RotatedFasterRCNN:
+    return RotatedFasterRCNN(
+        num_classes=3,
+        backbone_name="resnet18",
+        pretrained_backbone=False,
+        trainable_layers=5,
+        rpn_post_nms_top_n=128,
+        rpn_pre_nms_top_n=128,
+        inference_pre_nms_score_threshold=0.01,
+        final_nms_iou_threshold=0.3,
+        nms_class_agnostic=True,
+        final_nms_use_cpu=True,
+        max_detections_per_image=64,
+        roi_inference_top_class_only=True,
+    )
+
+
+def _finalize_kwargs_from_model(model: RotatedFasterRCNN) -> dict:
+    return {
+        "nms_class_agnostic": model.nms_class_agnostic,
+        "final_nms_iou_threshold": model.final_nms_iou_threshold,
+        "max_detections_per_image": model.max_detections_per_image,
+        "final_nms_use_cpu": model.final_nms_use_cpu,
+        "score_threshold": 0.05,
+        "per_class_score_threshold": None,
+        "class_id_to_name": {i + 1: f"class_{i}" for i in range(model.num_classes)},
+        "max_output_slots": model.max_detections_per_image,
+    }
+
+
+def test_pre_nms_finalize_matches_full_inference(tiny_faster_rcnn: RotatedFasterRCNN) -> None:
+    """Export path + postprocess should match faster_rcnn_inference (deterministic RPN)."""
+    model = tiny_faster_rcnn
+    model.eval()
+    h, w = 128, 128
+    torch.manual_seed(0)
+    x = torch.rand(1, 3, h, w, dtype=torch.float32)
+
+    with torch.no_grad():
+        full_out = faster_rcnn_inference(model, [x[0]], deterministic_rpn=True)[0]
+        wrap = RotatedFasterRCNNPreNmsExportWrapper(model, height=h, width=w, max_candidates=128)
+        boxes, scores, labels, count = wrap(x)
+        det_np, num = finalize_detections_numpy(
+            boxes.cpu().numpy(),
+            scores.cpu().numpy(),
+            labels.cpu().numpy(),
+            int(count.item()),
+            **_finalize_kwargs_from_model(model),
+        )
+
+    if full_out["scores"].numel() == 0:
+        assert num == 0
+        return
+
+    full_boxes = rboxes_to_tensor(full_out["rboxes"])
+    full_scores = full_out["scores"]
+    full_labels = full_out["labels"]
+
+    assert num == int(full_boxes.shape[0])
+    assert torch.allclose(
+        torch.from_numpy(det_np[:num, :5]),
+        full_boxes,
+        atol=1e-4,
+        rtol=1e-4,
+    )
+    assert torch.allclose(
+        torch.from_numpy(det_np[:num, 5]),
+        full_scores,
+        atol=1e-4,
+        rtol=1e-4,
+    )
+    assert torch.equal(
+        torch.from_numpy(det_np[:num, 6].astype(np.int64)),
+        full_labels,
+    )
+
+
+def test_faster_rcnn_pre_nms_onnx_checker(tiny_faster_rcnn: RotatedFasterRCNN) -> None:
+    pytest.importorskip("onnx")
+    import io
+
+    import onnx
+
+    model = tiny_faster_rcnn
+    model.eval()
+    h, w = 128, 128
+    wrap = RotatedFasterRCNNPreNmsExportWrapper(model, height=h, width=w, max_candidates=64)
+    x = torch.randn(1, 3, h, w, dtype=torch.float32)
+    buf = io.BytesIO()
+    torch.onnx.export(
+        wrap,
+        x,
+        buf,
+        input_names=["images"],
+        output_names=["pre_nms_boxes", "pre_nms_scores", "pre_nms_labels", "pre_nms_count"],
+        opset_version=17,
+        do_constant_folding=True,
+    )
+    buf.seek(0)
+    onnx.checker.check_model(onnx.load(buf))
+
+
+def test_faster_rcnn_pre_nms_onnx_ort_random_input(tiny_faster_rcnn: RotatedFasterRCNN) -> None:
+    """ORT must run on non-zero images (regression: pad_pre_nms used invalid Expand)."""
+    pytest.importorskip("onnx")
+    pytest.importorskip("onnxruntime")
+    import io
+
+    import numpy as np
+    import onnxruntime as ort
+
+    model = tiny_faster_rcnn
+    model.eval()
+    h, w = 128, 128
+    wrap = RotatedFasterRCNNPreNmsExportWrapper(model, height=h, width=w, max_candidates=64)
+    buf = io.BytesIO()
+    torch.onnx.export(
+        wrap,
+        torch.zeros(1, 3, h, w, dtype=torch.float32),
+        buf,
+        input_names=["images"],
+        output_names=["pre_nms_boxes", "pre_nms_scores", "pre_nms_labels", "pre_nms_count"],
+        opset_version=17,
+        do_constant_folding=True,
+    )
+    buf.seek(0)
+    sess = ort.InferenceSession(buf.getvalue(), providers=["CPUExecutionProvider"])
+    for seed in range(8):
+        torch.manual_seed(seed)
+        x = torch.rand(1, 3, h, w, dtype=torch.float32).numpy()
+        outs = sess.run(None, {"images": x})
+        assert outs[0].shape == (64, 5)
+        assert int(np.asarray(outs[3]).reshape(-1)[0]) >= 0
+
+
+@pytest.mark.skipif(
+    not _DEPLOY_CKPT.is_file() or not _DEPLOY_CONFIG.is_file(),
+    reason="deploy/app weights or config missing",
+)
+def test_deploy_checkpoint_pre_nms_finalize() -> None:
+    """Smoke parity on published deploy weights (slow; optional)."""
+    import json
+
+    from oriented_det.runtime.checkpoint import load_model_from_checkpoint
+
+    model, config, class_names = load_model_from_checkpoint(
+        str(_DEPLOY_CKPT), str(_DEPLOY_CONFIG), device="cpu"
+    )
+    assert isinstance(model, RotatedFasterRCNN)
+
+    meta = {
+        "production": json.loads(_DEPLOY_CONFIG.read_text()).get("production"),
+        "class_names": class_names or [],
+    }
+    kwargs = meta_to_finalize_kwargs(meta)
+    h, w = 256, 256
+    wrap = RotatedFasterRCNNPreNmsExportWrapper(
+        model, height=h, width=w, max_candidates=model.rpn_post_nms_top_n
+    )
+    x = torch.rand(1, 3, h, w, dtype=torch.float32)
+    with torch.no_grad():
+        boxes, scores, labels, count = wrap(x)
+        _, num = finalize_detections_numpy(
+            boxes.numpy(),
+            scores.numpy(),
+            labels.numpy(),
+            int(count.item()),
+            **kwargs,
+        )
+    assert num >= 0

export/tests/test_ort_runtime.py

@@ -0,0 +1,41 @@
+"""Tests for ORT device / provider resolution (no GPU required)."""
+
+from __future__ import annotations
+
+import pytest
+
+from export.ort_runtime import (
+    clear_ort_session_cache,
+    configure_ort_device,
+    get_ort_device,
+    ort_providers_for_device,
+    set_ort_device,
+)
+
+
+def test_cpu_providers():
+    set_ort_device("cpu")
+    assert ort_providers_for_device() == ["CPUExecutionProvider"]
+
+
+def test_unknown_device_raises():
+    with pytest.raises(ValueError, match="Unknown ORT device"):
+        ort_providers_for_device("tpu")
+
+
+def test_configure_returns_cpu_by_default():
+    clear_ort_session_cache()
+    set_ort_device(None)
+    providers = configure_ort_device("cpu")
+    assert providers == ["CPUExecutionProvider"]
+    assert get_ort_device() == "cpu"
+
+
+def test_cuda_without_gpu_raises():
+    import onnxruntime as ort
+
+    if "CUDAExecutionProvider" in ort.get_available_providers():
+        pytest.skip("CUDA EP available; skip missing-CUDA test")
+    set_ort_device("cuda")
+    with pytest.raises(RuntimeError, match="CUDAExecutionProvider"):
+        ort_providers_for_device()

export/tf_serving_model.py

@@ -0,0 +1,96 @@
+"""Keras-serializable Faster R-CNN detect pipeline (ORT core + Python rotated NMS)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import tensorflow as tf
+
+from export.postprocess import finalize_detections_numpy, ort_pre_nms_to_detections
+
+
+@tf.keras.utils.register_keras_serializable(package="oriented_det_export")
+class FasterRCNNDetectLayer(tf.keras.layers.Layer):
+    """Single end-to-end detect: ONNX Runtime pre-NMS + exact CPU rotated NMS."""
+
+    def __init__(
+        self,
+        onnx_path: str,
+        ort_output_names: List[str],
+        finalize_kwargs: Dict[str, Any],
+        max_output_slots: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.onnx_path = str(onnx_path)
+        self.ort_output_names = list(ort_output_names)
+        self.finalize_kwargs = dict(finalize_kwargs)
+        self.max_output_slots = int(max_output_slots)
+
+    def call(self, images: tf.Tensor) -> Dict[str, tf.Tensor]:
+        detections, num_detections = tf.numpy_function(
+            self._infer_numpy,
+            [images],
+            (tf.float32, tf.int32),
+        )
+        detections.set_shape([self.max_output_slots, 7])
+        num_detections.set_shape([])
+        return {"detections": detections, "num_detections": num_detections}
+
+    def _infer_numpy(self, images: np.ndarray) -> tuple[np.ndarray, np.int32]:
+        detections, num = ort_pre_nms_to_detections(
+            images,
+            self.onnx_path,
+            self.ort_output_names,
+            self.finalize_kwargs,
+        )
+        return detections.astype(np.float32), np.int32(num)
+
+    def get_config(self) -> Dict[str, Any]:
+        return {
+            "onnx_path": self.onnx_path,
+            "ort_output_names": self.ort_output_names,
+            "finalize_kwargs": self.finalize_kwargs,
+            "max_output_slots": self.max_output_slots,
+        }
+
+
+def build_keras_detect_model(
+    onnx_path: Path,
+    meta: Dict[str, Any],
+    finalize_kwargs: Dict[str, Any],
+) -> tf.keras.Model:
+    max_out = int(finalize_kwargs["max_output_slots"])
+    inputs = tf.keras.Input(shape=(3, None, None), batch_size=1, name="images")
+    # Fixed H,W enforced at call time; channel-first matches export convention.
+    out = FasterRCNNDetectLayer(
+        onnx_path=str(onnx_path.resolve()),
+        ort_output_names=list(meta.get("output_names") or []),
+        finalize_kwargs=finalize_kwargs,
+        max_output_slots=max_out,
+        name="faster_rcnn_detect",
+    )(inputs)
+    return tf.keras.Model(inputs=inputs, outputs=[out["detections"], out["num_detections"]])
+
+
+def save_keras_detect_bundle(
+    model: tf.keras.Model,
+    output_path: Path,
+    meta: Dict[str, Any],
+) -> Path:
+    """Save ``keras_model.keras`` + ``export_meta.json`` (load with :func:`load_keras_detect_model`)."""
+    output_path.mkdir(parents=True, exist_ok=True)
+    keras_path = output_path / "keras_model.keras"
+    model.save(keras_path)
+    (output_path / "export_meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
+    return keras_path
+
+
+def load_keras_detect_model(path: Path) -> tf.keras.Model:
+    return tf.keras.models.load_model(
+        path,
+        custom_objects={"FasterRCNNDetectLayer": FasterRCNNDetectLayer},
+    )

export/val_dataset.py

@@ -0,0 +1,116 @@
+"""Validation image enumeration (shared with tools/save_predictions.py semantics)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from oriented_det.data.airbus_playground import AirbusPlaygroundCSVDataset
+from oriented_det.data.dota import collect_dota_split_image_paths
+from oriented_det.train.config import TrainingExperimentConfig
+
+
+def collect_split_images(
+    config: TrainingExperimentConfig,
+    data_root: Path,
+    data_split: str = "val",
+    val_dir: Optional[Path] = None,
+    *,
+    filter_empty_gt: Optional[bool] = None,
+) -> Tuple[List[Path], Optional[Path], str]:
+    """Return image paths, label directory (DOTA), and dataset format string.
+
+    ``filter_empty_gt``: when ``None``, use ``config.dataset.filter_empty_gt`` (training).
+    ``tools/save_predictions.py`` (``make preds`` / ``make metrics``) always passes ``False``
+    so inference includes all tiles; training still uses the config flag.
+    """
+    dataset_format = getattr(getattr(config, "dataset", None), "format", None) or "dota"
+    data_root = Path(data_root)
+
+    if dataset_format == "airbus_playground":
+        ds_config = config.dataset
+        if not getattr(ds_config, "annotations_file", None) or not getattr(ds_config, "split_file", None):
+            raise ValueError(
+                "Airbus Playground format requires dataset.annotations_file and dataset.split_file."
+            )
+        if data_split not in ("train", "val"):
+            raise ValueError(f"Airbus dataset supports only train or val split, got {data_split!r}.")
+        airbus_dataset = AirbusPlaygroundCSVDataset(
+            data_root=data_root,
+            split=data_split,
+            annotations_file=ds_config.annotations_file,
+            split_file=ds_config.split_file,
+            val_split_id=getattr(ds_config, "val_split_id", 0),
+            difficult_strategy=ds_config.difficult_strategy,
+            allowed_classes=getattr(ds_config, "allowed_classes", None),
+            ignore_labels=getattr(ds_config, "ignore_labels", None) or [],
+            map_labels=getattr(ds_config, "map_labels", None) or {},
+        )
+        split_images = [Path(airbus_dataset[idx].image_path) for idx in range(len(airbus_dataset))]
+        return split_images, None, dataset_format
+
+    if getattr(config, "dataset", None):
+        if data_split == "val" and val_dir is not None:
+            tile_roots = [Path(val_dir)]
+        elif data_split == "val":
+            tile_roots = config.dataset.get_val_tile_roots()
+        elif data_split == "train":
+            tile_roots = config.dataset.get_train_tile_roots()
+        else:
+            tile_roots = [data_root / data_split]
+        same_folder = getattr(config.dataset, "same_folder", False)
+        difficult_strategy = getattr(config.dataset, "difficult_strategy", "drop")
+        allowed_classes = getattr(config.dataset, "allowed_classes", None)
+        ignore_labels = getattr(config.dataset, "ignore_labels", None)
+        if filter_empty_gt is None:
+            filter_empty_gt = bool(getattr(config.dataset, "filter_empty_gt", False))
+        split_images = collect_dota_split_image_paths(
+            tile_roots,
+            split=data_split,
+            same_folder=same_folder,
+            difficult_strategy=difficult_strategy,
+            allowed_classes=allowed_classes,
+            ignore_labels=ignore_labels,
+            filter_empty_gt=filter_empty_gt,
+            log_filter_empty_gt=filter_empty_gt,
+        )
+        if not split_images:
+            raise ValueError(f"No images found under DOTA tile roots: {tile_roots}")
+        label_dir = None
+        if len(tile_roots) == 1:
+            split_root = tile_roots[0]
+            split_tiles_dir = split_root / "tiles_1024"
+            if same_folder:
+                label_dir = split_root
+            elif (split_root / "labels").exists():
+                label_dir = split_root / "labels"
+            elif split_tiles_dir.exists() and (split_tiles_dir / "labels").exists():
+                label_dir = split_tiles_dir / "labels"
+            elif (split_root / "labelTxt").exists():
+                label_dir = split_root / "labelTxt"
+        return split_images, label_dir, dataset_format
+
+    split_root = data_root / data_split
+    split_tiles_dir = split_root / "tiles_1024"
+    same_folder = getattr(getattr(config, "dataset", None), "same_folder", False)
+
+    if same_folder:
+        image_dir = split_root
+        label_dir = split_root
+    elif split_tiles_dir.exists() and (split_tiles_dir / "images").exists():
+        image_dir = split_tiles_dir / "images"
+        label_dir = split_tiles_dir / "labels"
+    elif (split_root / "images").exists():
+        image_dir = split_root / "images"
+        label_dir = (
+            split_root / "labels"
+            if (split_root / "labels").exists()
+            else split_root / "labelTxt"
+        )
+    else:
+        raise ValueError(f"Could not find images under {split_tiles_dir} or {split_root}")
+
+    split_images = sorted(list(image_dir.glob("*.jpg")) + list(image_dir.glob("*.png")))
+    if not split_images:
+        raise ValueError(f"No images found in {image_dir}")
+    return split_images, label_dir, dataset_format

export/wrappers.py

@@ -0,0 +1,161 @@
+"""ONNX-friendly ``nn.Module`` wrappers around oriented-det PyTorch models.
+
+Export targets **tensor-in / tensor-out** subgraphs. Post-decode NMS and RBox
+assembly stay in Python or a separate TF/TFLite graph (see export/PARITY.md).
+"""
+
+from __future__ import annotations
+
+from typing import List, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from oriented_det.models.faster_rcnn_inference import faster_rcnn_inference_pre_nms_padded
+from oriented_det.models.oriented_rcnn import RotatedFasterRCNN
+from oriented_det.models.oriented_rpn import generate_oriented_anchors
+from oriented_det.models.rotated_retinanet import RotatedRetinaNet
+from oriented_det.models.utils import derive_fpn_strides_from_grid, extract_backbone_features
+
+
+def _ordered_fpn_values(features: object) -> Tuple[torch.Tensor, ...]:
+    """Turn backbone output into a tuple of tensors in a stable order."""
+    if isinstance(features, dict):
+        keys = []
+        for k in sorted(features.keys(), key=lambda x: str(x)):
+            ks = str(k)
+            if ks.isdigit() or ks.startswith("fpn"):
+                keys.append(k)
+        if not keys:
+            keys = list(features.keys())
+        return tuple(features[k] for k in keys)
+    if isinstance(features, (list, tuple)):
+        return tuple(features)
+    return (features,)  # type: ignore[return-value]
+
+
+class BackboneExportWrapper(nn.Module):
+    """Exports ``backbone(images)`` as a flat tuple of FPN tensors.
+
+    Works for any model that exposes a ``backbone`` attribute (ResNet+FPN).
+
+    Args:
+        backbone: Module accepting ``[B, 3, H, W]`` float tensor.
+    """
+
+    def __init__(self, backbone: nn.Module) -> None:
+        super().__init__()
+        self.backbone = backbone
+
+    def forward(self, images: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        feats = self.backbone(images)
+        return _ordered_fpn_values(feats)
+
+
+class RetinaNetBackboneHeadExportWrapper(nn.Module):
+    """Exports Rotated RetinaNet **backbone + classification/regression head** only.
+
+    Outputs alternate per FPN level: ``cls_lvl0``, ``bbox_lvl0``, ``cls_lvl1``, …
+    Classification tensors are sigmoid logits ``[B, A*K, H, W]``; bbox tensors are
+    ``[B, A*5, H, W]`` in model order. Decoding, anchor generation, thresholding,
+    and oriented NMS are **not** included (run in PyTorch inference or downstream TF).
+    """
+
+    def __init__(self, model: RotatedRetinaNet) -> None:
+        super().__init__()
+        self.backbone = model.backbone
+        self.extra_fpn_conv = model.extra_fpn_conv
+        self.head = model.head
+
+    def forward(self, images: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        # List of (C, H, W) — use unbind so batch size stays symbolic for ONNX.
+        img_list: Sequence[torch.Tensor] = torch.unbind(images, dim=0)
+        feat_list = extract_backbone_features(
+            self.backbone,
+            img_list,
+            use_checkpoint=False,
+            training=False,
+            include_pool_level=True,  # match RotatedRetinaNet.forward (P6 pool level)
+        )
+        if self.extra_fpn_conv is not None:
+            extra = F.relu(self.extra_fpn_conv(feat_list[-1]))
+            feat_list = list(feat_list) + [extra]
+        cls_list, bbox_list = self.head(feat_list)
+        out: List[torch.Tensor] = []
+        for c, b in zip(cls_list, bbox_list):
+            out.append(c)
+            out.append(b)
+        return tuple(out)
+
+
+class RotatedFasterRCNNPreNmsExportWrapper(nn.Module):
+    """Rotated Faster R-CNN through ROI decode; outputs padded pre-NMS tensors.
+
+    Input: ``images`` ``[1, 3, H, W]`` float32 in [0, 1] RGB.
+    Outputs: ``pre_nms_boxes``, ``pre_nms_scores``, ``pre_nms_labels``, ``pre_nms_count``.
+    Final rotated NMS and production score filters run in the TF SavedModel wrapper.
+    """
+
+    def __init__(
+        self,
+        model: RotatedFasterRCNN,
+        height: int,
+        width: int,
+        max_candidates: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self.height = int(height)
+        self.width = int(width)
+        self.max_candidates = int(max_candidates or model.rpn_post_nms_top_n)
+
+        with torch.no_grad():
+            dummy = torch.zeros(1, 3, self.height, self.width, dtype=torch.float32)
+            feat_list = extract_backbone_features(
+                model.backbone,
+                [dummy[0]],
+                use_checkpoint=False,
+                training=False,
+                include_pool_level=True,  # P6 for the RPN, matching RotatedFasterRCNN.forward
+            )
+            feature_map_sizes = [(f.shape[2], f.shape[3]) for f in feat_list]
+            fpn_strides_live = derive_fpn_strides_from_grid(
+                (self.height, self.width), feature_map_sizes
+            )
+            anchors = generate_oriented_anchors(
+                image_size=(self.height, self.width),
+                feature_map_sizes=feature_map_sizes,
+                anchor_scales=model.anchor_scales,
+                anchor_ratios=model.anchor_ratios,
+                anchor_angles=model.anchor_angles,
+                stride_per_level=fpn_strides_live,
+            )
+
+        self._num_anchor_levels = len(anchors)
+        for i, anchor_tensor in enumerate(anchors):
+            self.register_buffer(f"_anchor_{i}", anchor_tensor)
+        self._fpn_strides_list = [int(s) for s in fpn_strides_live]
+
+    def _anchor_list(self) -> List[torch.Tensor]:
+        return [getattr(self, f"_anchor_{i}") for i in range(self._num_anchor_levels)]
+
+    def forward(
+        self, images: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        return faster_rcnn_inference_pre_nms_padded(
+            self.model,
+            images,
+            self.max_candidates,
+            anchors=self._anchor_list(),
+            fpn_strides_live=self._fpn_strides_list,
+            deterministic_rpn=True,
+        )
+
+
+__all__ = [
+    "BackboneExportWrapper",
+    "RetinaNetBackboneHeadExportWrapper",
+    "RotatedFasterRCNNPreNmsExportWrapper",
+    "_ordered_fpn_values",
+]

oriented_det/__init__.py

@@ -0,0 +1,77 @@
+"""OrientedDet: rotated object detection for aerial imagery.
+
+Exports geometry (``RBox``, ``QBox``, ``Polygon``), ops (``iou``, ``nms``), data loaders,
+and detectors (``OrientedRCNN``, ``RotatedFasterRCNN``, ``RotatedRetinaNet``).
+Training configs: ``oriented_det.train.config.TrainingExperimentConfig``.
+"""
+
+from __future__ import annotations
+
+from .geometry import Polygon, QBox, RBox, transforms as geometry_transforms
+
+# Alias for convenience (geometry_transforms is the canonical name)
+transforms = geometry_transforms
+
+__all__ = [
+    "Polygon",
+    "QBox",
+    "RBox",
+    "transforms",
+    "geometry_transforms",
+    "iou",
+    "nms",
+    "DOTADataset",
+    "DOTASample",
+    "DOTAAnnotation",
+    "build_dota_loader",
+    "ImageTiler",
+    "compute_oriented_map",
+    "OrientedRCNN",
+    "RotatedFasterRCNN",
+    "RotatedRetinaNet",
+    "DeltaXYWHBBoxCoder",   # 4 params, RPN (MMRotate Rotated Faster R-CNN)
+    "DeltaXYWHAHBBoxCoder", # 5 params, ROI (MMRotate)
+    "MidpointOffsetCoder",  # 6 params, RPN (MMRotate Oriented R-CNN)
+]
+
+
+_LAZY_ATTRS = {
+    # ops
+    "iou": (".ops", "iou"),
+    "nms": (".ops", "nms"),
+    # data
+    "DOTADataset": (".data", "DOTADataset"),
+    "DOTASample": (".data", "DOTASample"),
+    "DOTAAnnotation": (".data", "DOTAAnnotation"),
+    "build_dota_loader": (".data", "build_dota_loader"),
+    "ImageTiler": (".data", "ImageTiler"),
+    "compute_oriented_map": (".data", "compute_oriented_map"),
+    # models
+    "OrientedRCNN": (".models", "OrientedRCNN"),
+    "RotatedFasterRCNN": (".models", "RotatedFasterRCNN"),
+    "RotatedRetinaNet": (".models", "RotatedRetinaNet"),
+    "DeltaXYWHBBoxCoder": (".models", "DeltaXYWHBBoxCoder"),
+    "DeltaXYWHAHBBoxCoder": (".models", "DeltaXYWHAHBBoxCoder"),
+    "MidpointOffsetCoder": (".models", "MidpointOffsetCoder"),
+}
+
+
+def __getattr__(name: str):
+    """Lazy import torch-heavy modules on first access.
+
+    This keeps lightweight utilities (e.g. `oriented_det.pretrained`) usable in
+    environments where importing PyTorch/torchvision kernels is undesirable.
+    """
+    if name not in _LAZY_ATTRS:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_path, attr = _LAZY_ATTRS[name]
+    import importlib
+
+    module = importlib.import_module(module_path, __name__)
+    value = getattr(module, attr)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals().keys()) | set(_LAZY_ATTRS.keys()))