team_code.py

#!/usr/bin/env python

# Edit this script to add your team's code. Some functions are *required*, but you can edit most parts of the required functions,
# change or remove non-required functions, and add your own functions.

################################################################################
#
# Optional libraries, functions, and variables. You can change or remove them.
#
################################################################################

import joblib
import os
import atexit
import builtins
import re
import sys

import numpy as np
from tqdm import tqdm

from helper_code import *
from src.pipeline.config import DEFAULT_CSV_PATH
from src.pipeline.features import get_feature_export_dir, get_or_create_record_feature_vector
from src.pipeline.training import (
    export_feature_views,
    predict_ensemble_labels,
    train_multimodal_ensemble,
)
################################################################################
# Path & Constant Configuration (Added for Robustness)
################################################################################

# Progress bar state for run_model (initialized lazily)
RUN_MODEL_PBAR = None
RUN_MODEL_PBAR_TOTAL = None
RUN_MODEL_EXPORT_STATE = None
ORIGINAL_PRINT = builtins.print
PRINT_FILTER_ACTIVE = False
RUN_PROGRESS_LINE_RE = re.compile(r'^-\s+\d+/\d+:\s')
def _close_run_model_pbar():
    global RUN_MODEL_PBAR
    if RUN_MODEL_PBAR is not None:
        RUN_MODEL_PBAR.close()
        RUN_MODEL_PBAR = None


def _flush_run_feature_exports():
    global RUN_MODEL_EXPORT_STATE

    if RUN_MODEL_EXPORT_STATE is None or not RUN_MODEL_EXPORT_STATE['metadata_rows']:
        return

    feature_exports = export_feature_views(
        RUN_MODEL_EXPORT_STATE['export_root'],
        'test',
        RUN_MODEL_EXPORT_STATE['metadata_rows'],
        np.asarray(RUN_MODEL_EXPORT_STATE['raw_features'], dtype=np.float32),
        RUN_MODEL_EXPORT_STATE['feature_names'],
        preprocessor=RUN_MODEL_EXPORT_STATE['preprocessor'],
    )

    if RUN_MODEL_EXPORT_STATE['verbose']:
        print(f"Raw features CSV: {feature_exports.get('raw')}")
        print(f"Preprocessed features CSV: {feature_exports.get('preprocessed')}")

    RUN_MODEL_EXPORT_STATE = None


def _install_run_print_filter():
    global PRINT_FILTER_ACTIVE
    if PRINT_FILTER_ACTIVE:
        return

    def _filtered_print(*args, **kwargs):
        message = kwargs.get('sep', ' ').join(str(a) for a in args) if args else ''
        if RUN_PROGRESS_LINE_RE.match(message):
            return
        return ORIGINAL_PRINT(*args, **kwargs)

    builtins.print = _filtered_print
    PRINT_FILTER_ACTIVE = True


def _restore_print():
    global PRINT_FILTER_ACTIVE
    if PRINT_FILTER_ACTIVE:
        builtins.print = ORIGINAL_PRINT
        PRINT_FILTER_ACTIVE = False


atexit.register(_close_run_model_pbar)
atexit.register(_flush_run_feature_exports)
atexit.register(_restore_print)


################################################################################
#
# Required functions. Edit these functions to add your code, but do not change the arguments for the functions.
#
################################################################################

# Train your models. This function is *required*. You should edit this function to add your code, but do *not* change the arguments
# of this function. If you do not train one of the models, then you can return None for the model.

# Train your model.
def train_model(data_folder, model_folder, verbose, csv_path=DEFAULT_CSV_PATH):
    # Find the data files.
    if verbose:
        print('Finding the Challenge data...')

    if verbose:
        print('Training the model on the data...')

    # Create a folder for the model if it does not already exist.
    os.makedirs(model_folder, exist_ok=True)

    model = train_multimodal_ensemble(
        data_folder,
        verbose,
        csv_path,
        export_folder=get_feature_export_dir(data_folder),
    )

    # Save the model.
    save_model(model_folder, model)

    feature_exports = model.get('feature_exports', {})
    if verbose and feature_exports:
        print(f"Raw features CSV: {feature_exports.get('raw')}")
        print(f"Preprocessed features CSV: {feature_exports.get('preprocessed')}")
        if feature_exports.get('selected'):
            print(f"Selected features CSV: {feature_exports.get('selected')}")

    if verbose:
        print('Done.')
        print()

# Load your trained models. This function is *required*. You should edit this function to add your code, but do *not* change the
# arguments of this function. If you do not train one of the models, then you can return None for the model.
def load_model(model_folder, verbose):
    if verbose:
        _install_run_print_filter()

    model_filename = os.path.join(model_folder, 'model.sav')
    model = joblib.load(model_filename)
    return model

# Run your trained model. This function is *required*. You should edit this function to add your code, but do *not* change the
# arguments of this function.
def run_model(model, record, data_folder, verbose):
    global RUN_MODEL_PBAR, RUN_MODEL_PBAR_TOTAL, RUN_MODEL_EXPORT_STATE

    # Load the model.
    model = model['model']

    # Extract identifiers from the record dictionary
    patient_id = record[HEADERS['bids_folder']]
    site_id    = record[HEADERS['site_id']]
    session_id = record[HEADERS['session_id']]

    # Initialize tqdm progress bar lazily so it advances across run_model calls.
    if verbose and RUN_MODEL_PBAR is None:
        patient_data_file = os.path.join(data_folder, DEMOGRAPHICS_FILE)
        try:
            RUN_MODEL_PBAR_TOTAL = len(find_patients(patient_data_file))
        except Exception:
            RUN_MODEL_PBAR_TOTAL = None

        RUN_MODEL_PBAR = tqdm(
            total=RUN_MODEL_PBAR_TOTAL,
            desc="Running Model",
            unit="record",
            leave=True,
            file=sys.stdout,
            delay=0.5,
            disable=not verbose
        )

    if RUN_MODEL_EXPORT_STATE is None:
        export_root = get_feature_export_dir(data_folder)
        RUN_MODEL_EXPORT_STATE = {
            'export_root': export_root,
            'metadata_rows': [],
            'raw_features': [],
            'feature_names': list(model.get('feature_names', [])),
            'preprocessor': model.get('preprocessor'),
            'verbose': verbose,
        }

    if verbose and RUN_MODEL_PBAR is not None:
        RUN_MODEL_PBAR.set_postfix({"patient": patient_id})

    # Load the patient data.
    patient_data_file = os.path.join(data_folder, DEMOGRAPHICS_FILE)
    patient_data = load_demographics(patient_data_file, patient_id, session_id)
    features = get_or_create_record_feature_vector(
        record,
        data_folder,
        patient_data,
        csv_path=DEFAULT_CSV_PATH,
        require_physiological_data=False,
    )
    features = np.asarray(features, dtype=np.float32).reshape(1, -1)

    RUN_MODEL_EXPORT_STATE['metadata_rows'].append({
        'patient_id': patient_id,
        'site_id': site_id,
        'session_id': session_id,
    })
    RUN_MODEL_EXPORT_STATE['raw_features'].append(features[0])

    # Get the model outputs.
    labels, probabilities = predict_ensemble_labels(model, features)
    binary_output = bool(int(labels[0]))
    probability_output = float(probabilities[0])

    if verbose and RUN_MODEL_PBAR is not None:
        RUN_MODEL_PBAR.update(1)
        if RUN_MODEL_PBAR_TOTAL is not None and RUN_MODEL_PBAR.n >= RUN_MODEL_PBAR_TOTAL:
            RUN_MODEL_PBAR.close()
            RUN_MODEL_PBAR = None

    if RUN_MODEL_PBAR_TOTAL is not None and RUN_MODEL_EXPORT_STATE is not None:
        if len(RUN_MODEL_EXPORT_STATE['metadata_rows']) >= RUN_MODEL_PBAR_TOTAL:
            _flush_run_feature_exports()

    return binary_output, probability_output

# Save your trained model.
def save_model(model_folder, model):
    d = {'model': model}
    filename = os.path.join(model_folder, 'model.sav')
    joblib.dump(d, filename, protocol=0)