Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,955 Bytes
a9a0ec2 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
# Copyright (c) Facebook, Inc. and its affiliates.
import numpy as np
import unittest
from typing import Dict
import torch
from detectron2.config import instantiate
from detectron2.structures import Boxes, Instances
class TestBaseHungarianTracker(unittest.TestCase):
def setUp(self):
self._img_size = np.array([600, 800])
self._prev_boxes = np.array(
[
[101, 101, 200, 200],
[301, 301, 450, 450],
]
).astype(np.float32)
self._prev_scores = np.array([0.9, 0.9])
self._prev_classes = np.array([1, 1])
self._prev_masks = np.ones((2, 600, 800)).astype("uint8")
self._curr_boxes = np.array(
[
[302, 303, 451, 452],
[101, 102, 201, 203],
]
).astype(np.float32)
self._curr_scores = np.array([0.95, 0.85])
self._curr_classes = np.array([1, 1])
self._curr_masks = np.ones((2, 600, 800)).astype("uint8")
self._prev_instances = {
"image_size": self._img_size,
"pred_boxes": self._prev_boxes,
"scores": self._prev_scores,
"pred_classes": self._prev_classes,
"pred_masks": self._prev_masks,
}
self._prev_instances = self._convertDictPredictionToInstance(self._prev_instances)
self._curr_instances = {
"image_size": self._img_size,
"pred_boxes": self._curr_boxes,
"scores": self._curr_scores,
"pred_classes": self._curr_classes,
"pred_masks": self._curr_masks,
}
self._curr_instances = self._convertDictPredictionToInstance(self._curr_instances)
self._max_num_instances = 200
self._max_lost_frame_count = 0
self._min_box_rel_dim = 0.02
self._min_instance_period = 1
self._track_iou_threshold = 0.5
def _convertDictPredictionToInstance(self, prediction: Dict) -> Instances:
"""
convert prediction from Dict to D2 Instances format
"""
res = Instances(
image_size=torch.IntTensor(prediction["image_size"]),
pred_boxes=Boxes(torch.FloatTensor(prediction["pred_boxes"])),
pred_masks=torch.IntTensor(prediction["pred_masks"]),
pred_classes=torch.IntTensor(prediction["pred_classes"]),
scores=torch.FloatTensor(prediction["scores"]),
)
return res
def test_init(self):
cfg = {
"_target_": "detectron2.tracking.hungarian_tracker.BaseHungarianTracker",
"video_height": self._img_size[0],
"video_width": self._img_size[1],
"max_num_instances": self._max_num_instances,
"max_lost_frame_count": self._max_lost_frame_count,
"min_box_rel_dim": self._min_box_rel_dim,
"min_instance_period": self._min_instance_period,
"track_iou_threshold": self._track_iou_threshold,
}
tracker = instantiate(cfg)
self.assertTrue(tracker._video_height == self._img_size[0])
def test_initialize_extra_fields(self):
cfg = {
"_target_": "detectron2.tracking.hungarian_tracker.BaseHungarianTracker",
"video_height": self._img_size[0],
"video_width": self._img_size[1],
"max_num_instances": self._max_num_instances,
"max_lost_frame_count": self._max_lost_frame_count,
"min_box_rel_dim": self._min_box_rel_dim,
"min_instance_period": self._min_instance_period,
"track_iou_threshold": self._track_iou_threshold,
}
tracker = instantiate(cfg)
instances = tracker._initialize_extra_fields(self._curr_instances)
self.assertTrue(instances.has("ID"))
self.assertTrue(instances.has("ID_period"))
self.assertTrue(instances.has("lost_frame_count"))
if __name__ == "__main__":
unittest.main()
|