yolo.TensorToBBoxConverter

Inheritance diagram of TensorToBBoxConverter — TensorToBBoxConverter inheritance diagram

class savant.converter.yolo.TensorToBBoxConverter(confidence_threshold=0.25, top_k=3000)

YOLO detector output to bbox converter.

__call__(*output_layers, model, roi)

Converts detector output layer tensor to bbox tensor.

Converter is suitable for PyTorch YOLOv5/v6/v7/v8 models. output_layers is assumed to consist of either one Nx(num_detected_classes+5) shape tensor, or one tensor of shape (num_detected_classes+4)xN, or 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]) – [top, left, 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

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.