text: #!/usr/bin/env python
  
  import numpy as np, os, sys, joblib
  from scipy.io import loadmat
  from sklearn.impute import SimpleImputer
  from sklearn.ensemble import RandomForestClassifier
  from sklearn.model_selection import GridSearchCV
  from get_12ECG_features import get_12ECG_features
  
  def train_12ECG_classifier(input_directory, output_directory):
      # Load data.
      print('Loading data...')
  
      header_files = []
      for f in os.listdir(input_directory):
          g = os.path.join(input_directory, f)
          if not f.lower().startswith('.') and f.lower().endswith('hea') and os.path.isfile(g):
              header_files.append(g)
  
      classes = get_classes(input_directory, header_files)
      num_classes = len(classes)
      num_files = len(header_files)
      recordings = list()
      headers = list()
  
      for i in range(num_files):
          recording, header = load_challenge_data(header_files[i])
          recordings.append(recording)
          headers.append(header)
  
      # Train model.
      print('Training model...')
  
      features = list()
      labels = list()
  
      for i in range(num_files):
          recording = recordings[i]
          header = headers[i]
  
          tmp = get_12ECG_features(recording, header)
          features.append(tmp)
  
          for l in header:
              if l.startswith('#Dx:'):
                  labels_act = np.zeros(num_classes)
                  arrs = l.strip().split(' ')
                  for arr in arrs[1].split(','):
                      class_index = classes.index(arr.rstrip()) # Only use first positive index
                      labels_act[class_index] = 1
          labels.append(labels_act)
  
      features = np.array(features)
      features = np.reshape(features,(num_files,1220))
      labels = np.array(labels)
  
      # Replace NaN values with mean values
      imputer=SimpleImputer().fit(features)
      features=imputer.transform(features)
  
      # Train the classifier
      # model = RandomForestClassifier().fit(features,labels)
      clf_rf = RandomForestClassifier(n_jobs=2)
  
      # Parameter range for grid search
      param_grid = {
      'n_estimators': [2, 4, 8, 16, 32, 64, 128, 256, 512],
      'max_features': ['auto', 'sqrt', 'log2'],
      'max_depth' : [2, 4, 8, 16, 32, 64, 128, 256, 512],
      'criterion' :['gini', 'entropy']
      }
      # Grid search
      CV_rf = GridSearchCV(estimator=clf_rf, param_grid=param_grid, cv= 4)
      CV_rf.fit(features, labels)
      CV_rf.best_params_
  
      model = RandomForestClassifier(max_features=CV_rf.best_params_['max_features'],
                                   n_estimators= CV_rf.best_params_['n_estimators'],
                                   max_depth=CV_rf.best_params_['max_depth'],
                                   criterion=CV_rf.best_params_['criterion'])
  
      # Retraining based on best parameters
      model.fit(features, labels)
  
      # Save model.
      print('Saving model...')
  
      final_model={'model':model, 'imputer':imputer,'classes':classes}
  
      filename = os.path.join(output_directory, 'finalized_model.sav')
      joblib.dump(final_model, filename, protocol=0)
  
  # Load challenge data.
  def load_challenge_data(header_file):
      with open(header_file, 'r') as f:
          header = f.readlines()
      mat_file = header_file.replace('.hea', '.mat')
      x = loadmat(mat_file)
      recording = np.asarray(x['val'], dtype=np.float64)
      return recording, header
  
  # Find unique classes.
  def get_classes(input_directory, filenames):
      classes = set()
      for filename in filenames:
          with open(filename, 'r') as f:
              for l in f:
                  if l.startswith('#Dx'):
                      tmp = l.split(': ')[1].split(',')
                      for c in tmp:
                          classes.add(c.strip())
      return sorted(classes)
  