yolo.TensorToBBoxConverter

TensorToBBoxConverter inheritance diagram

class savant.converter.yolo.TensorToBBoxConverter(confidence_threshold=0.25, nms_iou_threshold=0.0, top_k=3000, class_ids=None)

YOLOv4/v5/v6/v7/v8/v11 detector output to bbox converter.

__call__(*output_layers, model, roi)

Converts detector output layer tensor to bbox tensor.

Converter is suitable for PyTorch YOLOv4/v5/v6/v7/v8/v11 models. output_layers is assumed to consist of 1) 1 tensor of shape (num_detected_classes+4)xN, or 2) 1 tensor of shape Nx(num_detected_classes+4+1), or 3) 2 tensors of shapes Nx1x4 and Nx(num_detected_classes), or 4) 3 tensors of shapes Nx4, Nx(num_detected_classes), and Nx1, or 5) 4 tensors (after NMS) of shapes: 1, Nx4, N, N.

Parameters:

• output_layers (ndarray) – Output layer tensor

• model (ObjectModel) – Model definition, required parameters: input tensor shape, maintain_aspect_ratio

• roi (Tuple[float, float, float, float]) – [left, top, width, height] of the rectangle on which the model infers

Returns:

BBox tensor (class_id, confidence, xc, yc, width, height, [angle]) offset by roi upper left and scaled by roi width and height

Return type:

ndarray | None

tensor_format: TensorFormat = 0

Set to CuPy to get the output_layers tensors in the converter call on the GPU as a cupy.ndarray. Or set to NumPy to get tensors on the host as a numpy.ndarray.