yolo.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 models. output_layers is assumed to consist of either one Nx(num_detected_classes+5) shape tensor 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