k20hcmus/FishEye8K

1import threading2 3 4class Callbacks:5    """"6    Handles all registered callbacks for YOLOv5 Hooks7    """8 9    def __init__(self):10        # Define the available callbacks11        self._callbacks = {12            'on_pretrain_routine_start': [],13            'on_pretrain_routine_end': [],14            'on_train_start': [],15            'on_train_epoch_start': [],16            'on_train_batch_start': [],17            'optimizer_step': [],18            'on_before_zero_grad': [],19            'on_train_batch_end': [],20            'on_train_epoch_end': [],21            'on_val_start': [],22            'on_val_batch_start': [],23            'on_val_image_end': [],24            'on_val_batch_end': [],25            'on_val_end': [],26            'on_fit_epoch_end': [],  # fit = train + val27            'on_model_save': [],28            'on_train_end': [],29            'on_params_update': [],30            'teardown': [],}31        self.stop_training = False  # set True to interrupt training32 33    def register_action(self, hook, name='', callback=None):34        """35        Register a new action to a callback hook36 37        Args:38            hook: The callback hook name to register the action to39            name: The name of the action for later reference40            callback: The callback to fire41        """42        assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"43        assert callable(callback), f"callback '{callback}' is not callable"44        self._callbacks[hook].append({'name': name, 'callback': callback})45 46    def get_registered_actions(self, hook=None):47        """"48        Returns all the registered actions by callback hook49 50        Args:51            hook: The name of the hook to check, defaults to all52        """53        return self._callbacks[hook] if hook else self._callbacks54 55    def run(self, hook, *args, thread=False, **kwargs):56        """57        Loop through the registered actions and fire all callbacks on main thread58 59        Args:60            hook: The name of the hook to check, defaults to all61            args: Arguments to receive from YOLOv562            thread: (boolean) Run callbacks in daemon thread63            kwargs: Keyword Arguments to receive from YOLOv564        """65 66        assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"67        for logger in self._callbacks[hook]:68            if thread:69                threading.Thread(target=logger['callback'], args=args, kwargs=kwargs, daemon=True).start()70            else:71                logger['callback'](*args, **kwargs)72