File size: 1,106 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
#!/usr/bin/env python3
import numpy as np
from typing import List
from detectron2.structures import Instances
def create_prediction_pairs(
instances: Instances,
prev_instances: Instances,
iou_all: np.ndarray,
threshold: float = 0.5,
) -> List:
"""
Args:
instances: predictions from current frame
prev_instances: predictions from previous frame
iou_all: 2D numpy array containing iou for each bbox pair
threshold: below the threshold, doesn't consider the pair of bbox is valid
Return:
List of bbox pairs
"""
bbox_pairs = []
for i in range(len(instances)):
for j in range(len(prev_instances)):
if iou_all[i, j] < threshold:
continue
bbox_pairs.append(
{
"idx": i,
"prev_idx": j,
"prev_id": prev_instances.ID[j],
"IoU": iou_all[i, j],
"prev_period": prev_instances.ID_period[j],
}
)
return bbox_pairs
LARGE_COST_VALUE = 100000
|