text: # AUTOGENERATED! DO NOT EDIT! File to edit: 03_label_assistant.ipynb (unless otherwise specified).
  
  __all__ = ['LabelAssistant']
  
  # Cell
  import os
  import pandas as pd
  import winkie.dlc_importer as dlc_importer
  from .dlc_importer import DLCImporter
  from tqdm.notebook import tqdm
  
  class LabelAssistant:
      """Used to label data and generate labelled data.
  
      Please note that generating the labelled data will, as of now, also apply all pre-processing included in
      `dlc_importer`.
      """
  
      def __init__(self, behaviors, data_path, default_label='not_defined'):
          self.data_path = data_path
          self.behaviors = behaviors
          self.default_label = default_label
          self.files = os.listdir(data_path)
          self.labels = {'behavior': [], 'start': [], 'end': [], 'file': []}
          self.imp = DLCImporter()
  
      def addLabel(self, behavior, start_frame, end_frame, file_name):
          if behavior not in self.behaviors:
              raise ValueError(f"Unknown behaviour: {behavior}")
          if file_name not in self.files:
              raise FileNotFoundError(file_name)
  
          # There might be a user error of labellig frames that are not existent in the specified file.
          # We decide to ignore this potential user error for now, since we don't want to initially load the files.
  
          self.labels['behavior'].append(behavior)
          self.labels['start'].append(start_frame)
          self.labels['end'].append(end_frame)
          self.labels['file'].append(file_name)
  
      def apply_labels(self, origin_bodypart='body'):
          "Apply labels on all available data and execue preprocessing."
  
          def import_and_tag(file_name):
              "Note that this will also preprocess the `dlc_importer` coordinates."
              df = self.imp.import_hdf(file_name)
              # TODO: encapsulate sequence of steps into function in DLCImporter?
              df = dlc_importer.transform_to_relative(df, origin_bodypart)
              df = dlc_importer.add_middle_neck(df) # this is a bit opinionated
              df = dlc_importer.add_rotation(df)
              df = dlc_importer.apply_rotation(df)
              df['file_name'] = os.path.basename(file_name)
              df['frame'] = df.index
  
              return df
  
          filenames_to_import = list(filter(lambda x: x in self.labels['file'], self.files))
          files_to_import = [os.path.join(self.data_path, f) for f in filenames_to_import]
  
          df_list = [import_and_tag(f) for f in tqdm(files_to_import)]
          df_merged = pd.concat(df_list)
          df_merged['behavior'] = self.default_label
  
          for i, f in enumerate(tqdm(self.labels['file'])):
              file_name_select = df_merged['file_name'] == f
              frame_select = (self.labels['start'][i] <= df_merged['frame']) & (df_merged['frame'] <= self.labels['end'][i])
              df_merged.loc[file_name_select & frame_select, ['behavior']] = self.labels['behavior'][i]
  
          return df_merged
  
      def to_df(self):
          return pd.DataFrame(self.labels)