model.object_detection package

Submodules

model.object_detection.object_detector module

class model.object_detection.object_detector.ObjectDetector(name: str, model, inputs_spec: ImageSpec, predictions_spec: BoundingBoxSpec, preadapter: Callable[[...], Tuple[Tuple[Any, ...], Dict[str, Any]]] | None = None, postadapter: Callable[[Any], Any] | None = None, iou_threshold: float | None = None, score_threshold: float | None = None)

Bases: ArmoryModel, ModelProtocol

Wrapper around a model that produces object detection predictions.

This wrapper applies score-threshold filtering and non-maximum suppression to the model’s outputs if configured with a score threshold or an IOU threshold.

Example:

import armory.data
from armory.model.object_detection import ObjectDetector

# assuming `model` has been defined elsewhere
detector = ObjectDetector(
    name="My model",
    model=model,
    inputs_spec=armory.data.TorchImageSpec(
        dim=armory.data.ImageDimensions.CHW,
        scale=armory.data.Scale(
            dtype=armory.data.DataType.FLOAT,
            max=1.0,
        ),
    ),
    predictions_spec=armory.data.BoundingBoxSpec(
        format=armory.data.BBoxFormat.XYXY
    ),
)

predict(batch: ObjectDetectionBatch)

Invokes the wrapped model using the image inputs in the given batch and updates the object detection predictions in the batch.

If the wrapper has been configured with a score threshold or an IOU threshold, the predictions will be filtered before the batch predictions are updated.

Parameters:: batch (ObjectDetectionBatch) – Object detection batch

training: bool

model.object_detection.yolos_transformer module

Armory model wrapper for HuggingFace transformer YOLOS models.

class model.object_detection.yolos_transformer.YolosTransformer(name: str, model: YolosModel, image_processor: YolosImageProcessor, inputs_spec: ImageSpec | None = None, predictions_spec: BoundingBoxSpec | None = None, target_size: Tuple[int, int] = (512, 512), iou_threshold: float | None = None, score_threshold: float | None = None)

Bases: ObjectDetector

Model wrapper with pre-applied input and output adapters for HuggingFace transformer YOLOS models.

Example:

from transformers import AutoImageProcessor, AutoModelForObjectDetection
from armory.model.object_detection import YolosTransformer

model = AutoModelForObjectDetection.from_pretrained(CHECKPOINT)
processor = AutoImageProcessor.from_pretrained(CHECKPOINT)

transformer = YolosTransformer(
    name="My model",
    model=model,
    image_processor=processor,
)

DEFAULT_MEAN = (0.485, 0.456, 0.406)

DEFAULT_STD = (0.229, 0.224, 0.225)

training: bool

model.object_detection.yolov4_object_detector module

class model.object_detection.yolov4_object_detector.YoloV4ObjectDetector(name: str, model, inputs_spec: ImageSpec | None = None, predictions_spec: BoundingBoxSpec | None = None, iou_threshold: float | None = None, score_threshold: float | None = None)

Bases: ObjectDetector

Model wrapper with pre-applied output adapters for YOLOv4 models.

Example:

from armory.model.object_detection import Yolov4ObjectDetector

# assumes `model` has been created elsewhere
detector = YoloV4ObjectDetector(
    name="My model",
    model=model,
)

predict(batch: ObjectDetectionBatch)

Invokes the wrapped model using the image inputs in the given batch and updates the object detection predictions in the batch.

Non-maximum suppression processing is applied to the model’s outputs before the batch predictions are updated.

Model output / predictions format/example: outputs = self(inputs) Length of outputs: 2, type=list outputs[0] shape: torch.Size([1, 22743, 1, 4]) outputs[1] shape: torch.Size([1, 22743, 3])

predictions = _post_processing(outputs)[0] Length of predictions: 1, type=list[lists] predictions: [[-0.016737998, 0.051072836, 0.80007946, 0.96419203, 0.8180811, 0]] predictions[0]: [-0.016737998, 0.051072836, 0.80007946, 0.96419203, 0.8180811, 0]

Parameters:: batch (ObjectDetectionBatch) – Object detection batch

training: bool

model.object_detection.yolov5_object_detector module

class model.object_detection.yolov5_object_detector.YoloV5ObjectDetector(name: str, model, detection_model: DetectionModel | None = None, inputs_spec: ImageSpec | None = None, predictions_spec: BoundingBoxSpec | None = None, targets_spec: BoundingBoxSpec | None = None, iou_threshold: float | None = None, score_threshold: float | None = None, compute_loss: Callable[[Any, Any], Tuple[Any, Any]] | None = None, **kwargs)

Bases: ObjectDetector

Model wrapper with pre-applied output adapters for Ultralytics YOLOv5 models.

Example:

import yolov5
from armory.model.object_detection import Yolov5ObjectDetector

model = yolov5.load_model(CHECKPOINT)

detector = YoloV5ObjectDetector(
    name="My model",
    model=model,
)

forward(x, targets=None): Invokes the wrapped model. If in training and given targets, then the loss is computed and returned rather than the raw predictions. This is required when YoloV5ObjectDetector is used with ART and is wrapped by art.estimators.object_detection.PyTorchYolo.

loss(batch: ObjectDetectionBatch) → Tensor

Computes the loss for the given batch.

Parameters:: batch (ObjectDetectionBatch) – Object detection batch
Returns:: The total loss
Return type:: torch.Tensor

predict(batch: ObjectDetectionBatch)

Invokes the wrapped model using the image inputs in the given batch and updates the object detection predictions in the batch.

Non-maximum suppression processing is applied to the model’s outputs before the batch predictions are updated.

Parameters:: batch (ObjectDetectionBatch) – Object detection batch

training: bool

Module contents

This package contains model wrappers for object detection.