init commit of samurai

lib/test/tracker/basetracker.py

@@ -0,0 +1,89 @@
+import time
+
+import torch
+from _collections import OrderedDict
+
+from lib.train.data.processing_utils import transform_image_to_crop
+from lib.vis.visdom_cus import Visdom
+
+
+class BaseTracker:
+    """Base class for all trackers."""
+
+    def __init__(self, params):
+        self.params = params
+        self.visdom = None
+
+    def predicts_segmentation_mask(self):
+        return False
+
+    def initialize(self, image, info: dict) -> dict:
+        """Overload this function in your tracker. This should initialize the model."""
+        raise NotImplementedError
+
+    def track(self, image, info: dict = None) -> dict:
+        """Overload this function in your tracker. This should track in the frame and update the model."""
+        raise NotImplementedError
+
+    def visdom_draw_tracking(self, image, box, segmentation=None):
+        if isinstance(box, OrderedDict):
+            box = [v for k, v in box.items()]
+        else:
+            box = (box,)
+        if segmentation is None:
+            self.visdom.register((image, *box), 'Tracking', 1, 'Tracking')
+        else:
+            self.visdom.register((image, *box, segmentation), 'Tracking', 1, 'Tracking')
+
+    def transform_bbox_to_crop(self, box_in, resize_factor, device, box_extract=None, crop_type='template'):
+        # box_in: list [x1, y1, w, h], not normalized
+        # box_extract: same as box_in
+        # out bbox: Torch.tensor [1, 1, 4], x1y1wh, normalized
+        if crop_type == 'template':
+            crop_sz = torch.Tensor([self.params.template_size, self.params.template_size])
+        elif crop_type == 'search':
+            crop_sz = torch.Tensor([self.params.search_size, self.params.search_size])
+        else:
+            raise NotImplementedError
+
+        box_in = torch.tensor(box_in)
+        if box_extract is None:
+            box_extract = box_in
+        else:
+            box_extract = torch.tensor(box_extract)
+        template_bbox = transform_image_to_crop(box_in, box_extract, resize_factor, crop_sz, normalize=True)
+        template_bbox = template_bbox.view(1, 1, 4).to(device)
+
+        return template_bbox
+
+    def _init_visdom(self, visdom_info, debug):
+        visdom_info = {} if visdom_info is None else visdom_info
+        self.pause_mode = False
+        self.step = False
+        self.next_seq = False
+        if debug > 0 and visdom_info.get('use_visdom', True):
+            try:
+                self.visdom = Visdom(debug, {'handler': self._visdom_ui_handler, 'win_id': 'Tracking'},
+                                     visdom_info=visdom_info)
+
+                # # Show help
+                # help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \
+                #             'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \
+                #             'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \
+                #             'block list.'
+                # self.visdom.register(help_text, 'text', 1, 'Help')
+            except:
+                time.sleep(0.5)
+                print('!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n'
+                      '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!')
+
+    def _visdom_ui_handler(self, data):
+        if data['event_type'] == 'KeyPress':
+            if data['key'] == ' ':
+                self.pause_mode = not self.pause_mode
+
+            elif data['key'] == 'ArrowRight' and self.pause_mode:
+                self.step = True
+
+            elif data['key'] == 'n':
+                self.next_seq = True