Neuroimaging-Example/src/model.py

from .model_architecture import build_model
from .const_model import *
from .functions import *
import os
import mlflow

t1, t2, flair, t1ce, seg = modalities_path(base_path.replace('raw', 'processed'))
assert len(t1) == len(t2) == len(flair) == len(t1ce), 'Missing modalities in some patients'
assert len(t1) == len(t2) == len(flair) == len(t1ce) > 0, 'The lists of paths are empty'

num_samples: int = len(t1)

# TODO change it to a function the will not be redundant
data: np.array = np.empty((num_samples,) + input_shape, dtype=np.float32)
for index, item in enumerate(tqdm(zip(t1, t2, flair, t1ce), desc="Loading Data", total=len(t1))):
    data[index] = np.array([load_Nift2np(modal_path) for modal_path in item])

labels = load_raw_labels(seg, num_samples, input_shape, output_channels)  # TODO: Load the processed images.

mlflow.keras.autolog()

model = build_model(input_shape=input_shape, output_channels=output_channels)
model.fit(data, [labels, data], batch_size=batch_size, epochs=epochs)

model.save_weights(os.getcwd()+model_weights_path, overwrite=True)