def update_flag(key_press, current_flag, flags):
    """Handle key press from cv2.waitKey() for capturing frames

    :param key_press: output from `cv2.waitKey()`
    :param current_flag: value of 'flag' holding previous key press
    :param flags: dictionary mapping key presses to class labels
    :return: new flag value
    """
    if key_press < 0 or chr(key_press) not in flags.keys():
        return current_flag

    key_press = chr(key_press)
    for k in flags.keys():
        if k == key_press and k == current_flag:
            print(f'Stop capturing for {flags[k]}')
            return None
        elif k == key_press:
            print(f'Capturing for {flags[k]}')
            return k