TorchVisionODAdapter#
- class TorchVisionODAdapter[source]#
Bases:
Module
Adapter for interfacing with TorchVision’s object detection models.
The purpose of this adapter is: 1) to convert input BoxLists to dicts before feeding them into the model 2) to convert detections output by the model as dicts into BoxLists
Additionally, it automatically converts to/from 1-indexed class labels (which is what the TorchVision models expect).
- __init__(model: torch.nn.Module, ignored_output_inds: Sequence[int] = [0]) None [source]#
Constructor.
- Parameters
model (nn.Module) – A torchvision object detection model.
ignored_output_inds (Iterable[int], optional) – Class labels to exclude. Defaults to [0].
- Return type
None
Methods
__init__
(model[, ignored_output_inds])Constructor.
boxlist_to_model_input_dict
(boxlist)Convert BoxList to dict compatible w/ torchvision detection models.
forward
(input[, targets])Forward pass.
Convert model output dict to BoxList.
- __init__(model: torch.nn.Module, ignored_output_inds: Sequence[int] = [0]) None [source]#
Constructor.
- Parameters
model (nn.Module) – A torchvision object detection model.
ignored_output_inds (Iterable[int], optional) – Class labels to exclude. Defaults to [0].
- Return type
None
- boxlist_to_model_input_dict(boxlist: BoxList) dict [source]#
Convert BoxList to dict compatible w/ torchvision detection models.
Also, make class labels 1-indexed.
- forward(input: torch.Tensor, targets: Optional[Iterable[BoxList]] = None) Union[Dict[str, Any], List[BoxList]] [source]#
Forward pass.
- Parameters
input (Tensor[batch_size, in_channels, in_height, in_width]) – batch of images.
targets (Optional[Iterable[BoxList]], optional) – In training mode, should be Iterable[BoxList]], with each BoxList having a ‘class_ids’ field. In eval mode, should be None. Defaults to None.
- Returns
- In training mode,
returns a dict of losses. In eval mode, returns a list of BoxLists containing predicted boxes, class_ids, and scores. Further filtering based on score should be done before considering the prediction “final”.
- Return type