detector

The detection behind hbb2obb-detect: horizontal boxes for images that have none. See the guide for a worked example.

from pathlib import Path

import cv2

from hbb2obb.detector import detect_hbb, save_hbb_annotations

img_path = Path("data/images/04433.jpg")
boxes = detect_hbb(cv2.imread(str(img_path)), model="geotrax")  # (N, 6): class, xc, yc, w, h, conf
save_hbb_annotations(boxes, Path("data/labels_hbb"), img_path, scores=boxes[:, 5])

Detecting

detect_hbb

detect_hbb(img: ndarray, model: str = 'geotrax', imgsz: int = None, conf: float = None, iou: float = None, classes: Sequence[int] = None, class_map: Dict[int, int] = None, max_det: int = None, device: str = None, model_kwargs: Dict[str, Any] = None, models_dir: Optional[Path] = None) -> ndarray

Detect horizontal bounding boxes in one image.

Parameters:

Name	Type	Description	Default
img	ndarray	Image as read by cv2.imread (BGR)	required
model	str	Detector name, path or Hugging Face reference, as resolve_weights accepts	'geotrax'
imgsz	int	Inference resolution (default: the detector's own, 640 if it is not registered)	None
conf	float	Confidence threshold (default: the detector's own)	None
iou	float	NMS IoU threshold (default: the detector's own)	None
classes	Sequence[int]	Class IDs to keep, in the detector's own numbering (default: the detector's own)	None
class_map	Dict[int, int]	Remap the detector's class IDs onto yours; boxes of an unlisted class are dropped, so it selects and renumbers in one step	None
max_det	int	Maximum detections per image	None
device	str	Inference device, e.g. 'cpu', '0', 'mps' (default: Ultralytics picks)	None
model_kwargs	Dict[str, Any]	Any other Ultralytics predictor arguments, passed through unchecked; a key that is also given by name above raises ValueError	None
models_dir	Optional[Path]	Weights directory (default: $HBB2OBB_MODELS_DIR, else models)	None

Returns:

Type	Description
ndarray	An (N, 6) array of class, x_center, y_center, width, height, confidence, the coordinates in absolute pixels, ordered by descending confidence.

Source code in hbb2obb/detector.py

def detect_hbb(
    img: np.ndarray,
    model: str = "geotrax",
    imgsz: int = None,
    conf: float = None,
    iou: float = None,
    classes: Sequence[int] = None,
    class_map: Dict[int, int] = None,
    max_det: int = None,
    device: str = None,
    model_kwargs: Dict[str, Any] = None,
    models_dir: Optional[Path] = None,
) -> np.ndarray:
    """
    Detect horizontal bounding boxes in one image.

    Args:
        img: Image as read by ``cv2.imread`` (BGR)
        model: Detector name, path or Hugging Face reference, as ``resolve_weights`` accepts
        imgsz: Inference resolution (default: the detector's own, 640 if it is not registered)
        conf: Confidence threshold (default: the detector's own)
        iou: NMS IoU threshold (default: the detector's own)
        classes: Class IDs to keep, in the detector's own numbering (default: the detector's own)
        class_map: Remap the detector's class IDs onto yours; boxes of an unlisted class are
                   dropped, so it selects and renumbers in one step
        max_det: Maximum detections per image
        device: Inference device, e.g. 'cpu', '0', 'mps' (default: Ultralytics picks)
        model_kwargs: Any other Ultralytics predictor arguments, passed through unchecked; a key
                      that is also given by name above raises ``ValueError``
        models_dir: Weights directory (default: ``$HBB2OBB_MODELS_DIR``, else ``models``)

    Returns:
        An (N, 6) array of ``class, x_center, y_center, width, height, confidence``, the
        coordinates in absolute pixels, ordered by descending confidence.
    """
    kwargs = predictor_settings(model, imgsz, conf, iou, classes, max_det, device, model_kwargs)
    detector = load_detector(model, models_dir)

    result = detector(img, verbose=False, **kwargs)[0]
    boxes = result.boxes
    if boxes is None or len(boxes) == 0:
        return np.empty((0, 6))

    rows = np.column_stack(
        [
            boxes.cls.cpu().numpy().reshape(-1, 1),
            boxes.xywh.cpu().numpy(),
            boxes.conf.cpu().numpy().reshape(-1, 1),
        ]
    )
    return apply_class_map(rows, class_map)

save_hbb_annotations

save_hbb_annotations(hbb_annotations: ndarray, hbb_dir: Path, img_path: Path, scores: Union[Sequence[float], ndarray] = None, precision: int = 2, normalize: bool = False, img_shape: Tuple[int, int] = None) -> Path

Save HBB annotations to a text file, the format hbb2obb reads back.

Each line is class x_center y_center width height, in absolute pixels unless normalize is set. When scores is provided, the confidence is appended as a 6th field. Trailing zeros are trimmed, so a whole pixel is written as such rather than padded to the full precision.

Parameters:

Name	Type	Description	Default
hbb_annotations	ndarray	(N, 5) or (N, 6) array; a 6th column is ignored in favour of scores	required
hbb_dir	Path	Directory to write into (default: <img_path>/../../labels_hbb)	required
img_path	Path	Image the annotations belong to, which names the output file	required
scores	Union[Sequence[float], ndarray]	Per-box confidence, written as the trailing field	None
precision	int	Decimal places for the coordinates	2
normalize	bool	Write coordinates relative to the frame size instead of in pixels	False
img_shape	Tuple[int, int]	(width, height) of the frame, required when normalize is set	None

Returns:

Type	Description
Path	The path written.

Source code in hbb2obb/detector.py

def save_hbb_annotations(
    hbb_annotations: np.ndarray,
    hbb_dir: Path,
    img_path: Path,
    scores: Union[Sequence[float], np.ndarray] = None,
    precision: int = 2,
    normalize: bool = False,
    img_shape: Tuple[int, int] = None,
) -> Path:
    """
    Save HBB annotations to a text file, the format ``hbb2obb`` reads back.

    Each line is ``class x_center y_center width height``, in absolute pixels unless
    ``normalize`` is set. When ``scores`` is provided, the confidence is appended as a 6th
    field. Trailing zeros are trimmed, so a whole pixel is written as such rather than padded
    to the full precision.

    Args:
        hbb_annotations: (N, 5) or (N, 6) array; a 6th column is ignored in favour of ``scores``
        hbb_dir: Directory to write into (default: ``<img_path>/../../labels_hbb``)
        img_path: Image the annotations belong to, which names the output file
        scores: Per-box confidence, written as the trailing field
        precision: Decimal places for the coordinates
        normalize: Write coordinates relative to the frame size instead of in pixels
        img_shape: (width, height) of the frame, required when ``normalize`` is set

    Returns:
        The path written.
    """
    from hbb2obb.converter import _trim, format_annotation_line, resolve_output_dir

    if normalize and not img_shape:
        raise ValueError("img_shape is required to write normalized coordinates")

    hbb_dir = resolve_output_dir(hbb_dir, img_path, "labels_hbb")
    save_filepath = hbb_dir / (img_path.stem + ".txt")
    scale = np.array([*img_shape, *img_shape], dtype=float) if normalize else np.ones(4)
    scores = None if scores is None else list(scores)

    with open(save_filepath, "w", encoding="utf-8") as f:
        for i, box in enumerate(hbb_annotations):
            coords = np.asarray(box[1:5], dtype=float) / scale
            fields = f"{int(box[0])} " + " ".join(_trim(value, precision) for value in coords)
            f.write(format_annotation_line(fields, scores, i) + "\n")
    return save_filepath

load_detector

load_detector(model: str = 'geotrax', models_dir: Optional[Path] = None)

Load a detector by name or path, reusing a cached instance when one is already loaded.

Source code in hbb2obb/detector.py

def load_detector(model: str = "geotrax", models_dir: Optional[Path] = None):
    """Load a detector by name or path, reusing a cached instance when one is already loaded."""
    from ultralytics import YOLO

    weights = resolve_weights(model, models_dir)
    key = str(weights)
    if key not in _MODEL_CACHE:
        _MODEL_CACHE[key] = YOLO(weights)
    return _MODEL_CACHE[key]

resolve_weights

resolve_weights(model: str, models_dir: Optional[Path] = None) -> Path

Resolve a --model value to a weights file on disk, downloading it if it is not there yet.

Accepted, in this order: a name in SUPPORTED_DETECTORS; a path to an existing file; a Hugging Face file, as a link or as <user>/<repo>/<path>.pt; any other http(s) link to a checkpoint; and an Ultralytics model name such as yolo11s.pt. Downloads land in the models directory beside the SAM checkpoints (models_dir, else $HBB2OBB_MODELS_DIR, else models) under the file's own name, and Ultralytics names are left for Ultralytics itself to fetch on first use. A path-like value that exists nowhere is an error rather than a guess.

Source code in hbb2obb/detector.py

def resolve_weights(model: str, models_dir: Optional[Path] = None) -> Path:
    """
    Resolve a ``--model`` value to a weights file on disk, downloading it if it is not there yet.

    Accepted, in this order: a name in ``SUPPORTED_DETECTORS``; a path to an existing file; a
    Hugging Face file, as a link or as ``<user>/<repo>/<path>.pt``; any other ``http(s)`` link to a
    checkpoint; and an Ultralytics model name such as ``yolo11s.pt``. Downloads land in the models
    directory beside the SAM checkpoints (``models_dir``, else ``$HBB2OBB_MODELS_DIR``, else
    ``models``) under the file's own name, and Ultralytics names are left for Ultralytics itself to
    fetch on first use. A path-like value that exists nowhere is an error rather than a guess.
    """
    weights_dir = resolve_models_dir(models_dir)
    path = Path(model)
    if model not in SUPPORTED_DETECTORS and path.is_file():
        return path

    found = hub_file(model)
    if found is not None:
        return download_weights(found.url, weights_dir / PurePosixPath(found.path).name)

    if _is_url(model):
        name = PurePosixPath(urlparse(model).path).name
        if not name:
            raise ValueError(f"Cannot tell which file {model} names; link to the checkpoint itself")
        return download_weights(model, weights_dir / name)

    if "/" in model or "\\" in model:
        raise FileNotFoundError(
            f"{model} is not an existing file, a Hugging Face reference (<user>/<repo>/<file>.pt) or a link"
        )
    return weights_dir / (model if model.endswith(".pt") else f"{model}.pt")

Classes

parse_class_map

parse_class_map(spec: str) -> Dict[int, int]

Parse a --class_map string such as '2=0,5=1,7=2' into {2: 0, 5: 1, 7: 2}.

Source code in hbb2obb/detector.py

def parse_class_map(spec: str) -> Dict[int, int]:
    """Parse a ``--class_map`` string such as ``'2=0,5=1,7=2'`` into ``{2: 0, 5: 1, 7: 2}``."""
    if not spec:
        return {}

    mapping = {}
    for pair in spec.split(","):
        source, sep, target = pair.partition("=")
        if not sep:
            raise ValueError(f"Invalid class map entry '{pair}', expected 'source=target'")
        mapping[int(source)] = int(target)
    return mapping

apply_class_map

apply_class_map(rows: ndarray, class_map: Dict[int, int] = None) -> ndarray

Renumber the class column of detection rows, dropping the classes the map does not name.

A detector trained on COCO calls a car 2 and a bus 5; {2: 0, 5: 1} turns those into the 0 and 1 of a four-class vehicle label map and discards everything else.

Source code in hbb2obb/detector.py

def apply_class_map(rows: np.ndarray, class_map: Dict[int, int] = None) -> np.ndarray:
    """
    Renumber the class column of detection rows, dropping the classes the map does not name.

    A detector trained on COCO calls a car 2 and a bus 5; ``{2: 0, 5: 1}`` turns those into the
    0 and 1 of a four-class vehicle label map and discards everything else.
    """
    if not class_map or len(rows) == 0:
        return rows

    kept = np.array([int(cls) in class_map for cls in rows[:, 0]], dtype=bool)
    rows = rows[kept].copy()
    for i, cls in enumerate(rows[:, 0]):
        rows[i, 0] = class_map[int(cls)]
    return rows

Merging with hand-drawn boxes

merge_detections

merge_detections(manual: ndarray, detected: ndarray, iou_threshold: float = 0.5) -> MergeReport

Attach detector confidences to hand-drawn boxes without touching their geometry.

The hand-drawn set is authoritative: it stays as it is, in its own order, and each box only takes the confidence of the detection that overlaps it most. A box the detector missed gets 1.0, because a box somebody drew by hand is not less certain than one a model proposed, it is certain by construction. Detections that back no hand-drawn box are reported rather than added, since whether they are real objects the annotator missed is a judgement call.

Matching is greedy on IoU and class-agnostic, so a vehicle the detector called a truck and the annotator called a car still matches; the disagreement is reported in conflicts.

Parameters:

Name	Type	Description	Default
manual	ndarray	(M, 5) hand-drawn boxes as class, x_center, y_center, width, height	required
detected	ndarray	(N, 6) detections as class, x_center, y_center, width, height, confidence	required
iou_threshold	float	Minimum IoU for a detection to be considered the same object	0.5

Returns:

Type	Description
MergeReport	A MergeReport whose scores has one confidence per hand-drawn box, in order.

Source code in hbb2obb/detector.py

def merge_detections(manual: np.ndarray, detected: np.ndarray, iou_threshold: float = 0.5) -> MergeReport:
    """
    Attach detector confidences to hand-drawn boxes without touching their geometry.

    The hand-drawn set is authoritative: it stays as it is, in its own order, and each box only
    takes the confidence of the detection that overlaps it most. A box the detector missed gets
    1.0, because a box somebody drew by hand is not less certain than one a model proposed, it
    is certain by construction. Detections that back no hand-drawn box are reported rather than
    added, since whether they are real objects the annotator missed is a judgement call.

    Matching is greedy on IoU and class-agnostic, so a vehicle the detector called a truck and
    the annotator called a car still matches; the disagreement is reported in ``conflicts``.

    Args:
        manual: (M, 5) hand-drawn boxes as ``class, x_center, y_center, width, height``
        detected: (N, 6) detections as ``class, x_center, y_center, width, height, confidence``
        iou_threshold: Minimum IoU for a detection to be considered the same object

    Returns:
        A ``MergeReport`` whose ``scores`` has one confidence per hand-drawn box, in order.
    """
    report = MergeReport(scores=np.ones(len(manual)), extra=list(range(len(detected))))
    if len(manual) == 0 or len(detected) == 0:
        report.missed = list(range(len(manual)))
        return report

    ious = box_iou(xywh_to_xyxy(manual[:, 1:5]), xywh_to_xyxy(detected[:, 1:5]))
    taken_manual, taken_detected = set(), set()
    order = np.argsort(ious, axis=None)[::-1]
    for flat in order:
        i, j = np.unravel_index(flat, ious.shape)
        if ious[i, j] < iou_threshold:
            break
        if i in taken_manual or j in taken_detected:
            continue
        taken_manual.add(int(i))
        taken_detected.add(int(j))
        report.scores[i] = detected[j, 5]
        if int(manual[i, 0]) != int(detected[j, 0]):
            report.conflicts.append((int(i), int(j)))

    report.matched = sorted(taken_manual)
    report.missed = [i for i in range(len(manual)) if i not in taken_manual]
    report.extra = [j for j in range(len(detected)) if j not in taken_detected]
    return report

MergeReport dataclass

What merge_detections found, so that a caller can report it or review the leftovers.

Source code in hbb2obb/detector.py

@dataclass
class MergeReport:
    """What ``merge_detections`` found, so that a caller can report it or review the leftovers."""

    scores: np.ndarray = field(default_factory=lambda: np.empty((0,)))
    matched: List[int] = field(default_factory=list)
    missed: List[int] = field(default_factory=list)
    extra: List[int] = field(default_factory=list)
    conflicts: List[Tuple[int, int]] = field(default_factory=list)