Reference for ultralytics/utils/tal.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/tal.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.utils.tal.TaskAlignedAssigner
Bases: Module
A task-aligned assigner for object detection.
This class assigns ground-truth (gt) objects to anchors based on the task-aligned metric, which combines both classification and localization information.
Attributes:
Name | Type | Description |
---|---|---|
topk |
int
|
The number of top candidates to consider. |
num_classes |
int
|
The number of object classes. |
alpha |
float
|
The alpha parameter for the classification component of the task-aligned metric. |
beta |
float
|
The beta parameter for the localization component of the task-aligned metric. |
eps |
float
|
A small value to prevent division by zero. |
Source code in ultralytics/utils/tal.py
forward
Compute the task-aligned assignment. Reference code is available at https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/assigner/tal_assigner.py.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
pd_scores |
Tensor
|
shape(bs, num_total_anchors, num_classes) |
required |
pd_bboxes |
Tensor
|
shape(bs, num_total_anchors, 4) |
required |
anc_points |
Tensor
|
shape(num_total_anchors, 2) |
required |
gt_labels |
Tensor
|
shape(bs, n_max_boxes, 1) |
required |
gt_bboxes |
Tensor
|
shape(bs, n_max_boxes, 4) |
required |
mask_gt |
Tensor
|
shape(bs, n_max_boxes, 1) |
required |
Returns:
Name | Type | Description |
---|---|---|
target_labels |
Tensor
|
shape(bs, num_total_anchors) |
target_bboxes |
Tensor
|
shape(bs, num_total_anchors, 4) |
target_scores |
Tensor
|
shape(bs, num_total_anchors, num_classes) |
fg_mask |
Tensor
|
shape(bs, num_total_anchors) |
target_gt_idx |
Tensor
|
shape(bs, num_total_anchors) |
Source code in ultralytics/utils/tal.py
get_box_metrics
Compute alignment metric given predicted and ground truth bounding boxes.
Source code in ultralytics/utils/tal.py
get_pos_mask
Get in_gts mask, (b, max_num_obj, h*w).
Source code in ultralytics/utils/tal.py
get_targets
Compute target labels, target bounding boxes, and target scores for the positive anchor points.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
gt_labels |
Tensor
|
Ground truth labels of shape (b, max_num_obj, 1), where b is the batch size and max_num_obj is the maximum number of objects. |
required |
gt_bboxes |
Tensor
|
Ground truth bounding boxes of shape (b, max_num_obj, 4). |
required |
target_gt_idx |
Tensor
|
Indices of the assigned ground truth objects for positive anchor points, with shape (b, hw), where hw is the total number of anchor points. |
required |
fg_mask |
Tensor
|
A boolean tensor of shape (b, h*w) indicating the positive (foreground) anchor points. |
required |
Returns:
Type | Description |
---|---|
Tuple[Tensor, Tensor, Tensor]
|
A tuple containing the following tensors: - target_labels (Tensor): Shape (b, hw), containing the target labels for positive anchor points. - target_bboxes (Tensor): Shape (b, hw, 4), containing the target bounding boxes for positive anchor points. - target_scores (Tensor): Shape (b, h*w, num_classes), containing the target scores for positive anchor points, where num_classes is the number of object classes. |
Source code in ultralytics/utils/tal.py
iou_calculation
select_candidates_in_gts
staticmethod
Select the positive anchor center in gt.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
xy_centers |
Tensor
|
shape(h*w, 2) |
required |
gt_bboxes |
Tensor
|
shape(b, n_boxes, 4) |
required |
Returns:
Type | Description |
---|---|
Tensor
|
shape(b, n_boxes, h*w) |
Source code in ultralytics/utils/tal.py
select_highest_overlaps
staticmethod
If an anchor box is assigned to multiple gts, the one with the highest IoU will be selected.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
mask_pos |
Tensor
|
shape(b, n_max_boxes, h*w) |
required |
overlaps |
Tensor
|
shape(b, n_max_boxes, h*w) |
required |
Returns:
Name | Type | Description |
---|---|---|
target_gt_idx |
Tensor
|
shape(b, h*w) |
fg_mask |
Tensor
|
shape(b, h*w) |
mask_pos |
Tensor
|
shape(b, n_max_boxes, h*w) |
Source code in ultralytics/utils/tal.py
select_topk_candidates
Select the top-k candidates based on the given metrics.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
metrics |
Tensor
|
A tensor of shape (b, max_num_obj, hw), where b is the batch size, max_num_obj is the maximum number of objects, and hw represents the total number of anchor points. |
required |
largest |
bool
|
If True, select the largest values; otherwise, select the smallest values. |
True
|
topk_mask |
Tensor
|
An optional boolean tensor of shape (b, max_num_obj, topk), where topk is the number of top candidates to consider. If not provided, the top-k values are automatically computed based on the given metrics. |
None
|
Returns:
Type | Description |
---|---|
Tensor
|
A tensor of shape (b, max_num_obj, h*w) containing the selected top-k candidates. |
Source code in ultralytics/utils/tal.py
ultralytics.utils.tal.RotatedTaskAlignedAssigner
Bases: TaskAlignedAssigner
Assigns ground-truth objects to rotated bounding boxes using a task-aligned metric.
Source code in ultralytics/utils/tal.py
iou_calculation
select_candidates_in_gts
staticmethod
Select the positive anchor center in gt for rotated bounding boxes.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
xy_centers |
Tensor
|
shape(h*w, 2) |
required |
gt_bboxes |
Tensor
|
shape(b, n_boxes, 5) |
required |
Returns:
Type | Description |
---|---|
Tensor
|
shape(b, n_boxes, h*w) |
Source code in ultralytics/utils/tal.py
ultralytics.utils.tal.make_anchors
Generate anchors from features.
Source code in ultralytics/utils/tal.py
ultralytics.utils.tal.dist2bbox
Transform distance(ltrb) to box(xywh or xyxy).
Source code in ultralytics/utils/tal.py
ultralytics.utils.tal.bbox2dist
Transform bbox(xyxy) to dist(ltrb).
ultralytics.utils.tal.dist2rbox
Decode predicted object bounding box coordinates from anchor points and distribution.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
pred_dist |
Tensor
|
Predicted rotated distance, (bs, h*w, 4). |
required |
pred_angle |
Tensor
|
Predicted angle, (bs, h*w, 1). |
required |
anchor_points |
Tensor
|
Anchor points, (h*w, 2). |
required |
Returns: (torch.Tensor): Predicted rotated bounding boxes, (bs, h*w, 4).