def run_model(model_directory, data_directory, output_directory): batch_size = 1 # Load model. print('Loading models...') twelve_lead_model = load_twelve_lead_model(model_directory) six_lead_model = load_six_lead_model(model_directory) three_lead_model = load_three_lead_model(model_directory) two_lead_model = load_two_lead_model(model_directory) # Find header and recording files. print('Finding header and recording files...') header_files, recording_files = find_challenge_files(data_directory) num_recordings = len(recording_files) if not num_recordings: raise Exception('No data was provided.') # Create a folder for the outputs if it does not already exist. if not os.path.isdir(output_directory): os.mkdir(output_directory) # Run model for each recording. print('Running model...') for i in range(num_recordings): print(' {}/{}...'.format(i + 1, num_recordings)) # Load header and recording. header = load_header(header_files[i]) recording = load_recording(recording_files[i]) leads = get_leads(header) # Apply model to recording. if all(lead in leads for lead in twelve_leads): classes, labels, probabilities = run_twelve_lead_model( twelve_lead_model, header, recording) elif all(lead in leads for lead in six_leads): classes, labels, probabilities = run_six_lead_model( six_lead_model, header, recording) elif all(lead in leads for lead in three_leads): classes, labels, probabilities = run_three_lead_model( three_lead_model, header, recording) elif all(lead in leads for lead in two_leads): classes, labels, probabilities = run_two_lead_model( two_lead_model, header, recording) else: raise NotImplementedError( 'No model is implemented for the lead set {}.'.format( ', '.join(leads))) # Save model outputs. head, tail = os.path.split(header_files[i]) root, extension = os.path.splitext(tail) output_file = os.path.join(output_directory, root + '.csv') save_outputs(output_file, classes, labels, probabilities) print('Done.')
def test_model(model_directory, data_directory, output_directory): # Load model. print('Loading models...') twelve_lead_model = load_twelve_lead_model(model_directory) six_lead_model = load_six_lead_model(model_directory) three_lead_model = load_three_lead_model(model_directory) two_lead_model = load_two_lead_model(model_directory) # Find header and recording files. print('Finding header and recording files...') header_files, recording_files = find_challenge_files(data_directory) num_recordings = len(recording_files) if not num_recordings: raise Exception('No data was provided.') # Create a folder for the outputs if it does not already exist. if not os.path.isdir(output_directory): os.mkdir(output_directory) # Run model for each recording. print('Running model...') run_leads_model(twelve_lead_model, twelve_leads) run_leads_model(six_lead_model, six_leads) run_leads_model(three_lead_model, three_leads) run_leads_model(two_lead_model, two_leads)