added code

code/utils/inference.py

@@ -0,0 +1,213 @@
+import os
+
+import lmdb
+import numpy as np
+import torch
+from torch import nn
+
+from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature
+from utils.data_utils import get_collated_batch_for_key
+from utils.dataset_utils import get_dataset_path
+from utils.model_utils import get_model_config_from_file
+from utils.training_utils import get_training_config_from_file, get_data_ids, find_max_batch_size
+
+
+def load_model(model_dir: str,
+               results_base_dir: str,
+               lmdb_base_dir: str,
+               device: str = "cuda") -> tuple:
+    """
+    Load the model from the model directory.
+    :param model_dir: model directory
+    :param results_base_dir: results base directory
+    :param lmdb_base_dir: LMDB base directory
+    :param device: device to use
+    :return: model configuration, training configuration, dataset directory
+    """
+
+    training_id = None
+    model_configuration = get_model_config_from_file(model_dir, results_base_dir, lmdb_base_dir)
+    feature_config = model_configuration["feature_config"]
+    dataset_dir = get_dataset_path(lmdb_base_dir, feature_config["feature_set_name"])
+
+    # get list of trainings
+    training_base_dir = os.path.join(model_dir, "trainings")
+    # if no training id is given, use the latest training
+    if training_id is None:
+        training_dirs = os.listdir(training_base_dir)
+        training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
+                         os.path.isdir(os.path.join(training_base_dir, dir))]
+        training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
+        training_dir = training_dirs[0]
+    else:
+        training_dir = os.path.join(training_base_dir, training_id)
+
+    training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
+
+    limit = None
+    train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, limit)
+
+    dataset_dir = f"{lmdb_base_dir}/{feature_config['feature_set_name']}"
+    env = lmdb.open(dataset_dir, readonly=True)
+    load_fn = model_configuration["model_load_fn"]
+    model = load_fn(model_configuration, training_configuration, test_ids[0], device, lmdb_env=env)
+
+    # estimate batch size
+    sample_batch = get_collated_batch_for_key(train_ids[0], model_configuration, lmdb_env=env)
+    sample_input = sample_batch[0][0]  # get the first item in the batch
+    input_shape = sample_input.shape
+    # send model to proper device
+    model.to(device)
+    batch_size = find_max_batch_size(model,
+                                     input_shape,
+                                     device=device)
+    training_configuration["batch_size"] = batch_size
+    return model, model_configuration, training_configuration, train_ids, val_ids, test_ids
+
+
+def get_training(model_configuration: dict,
+                 results_base_dir: str,
+                 training_id: str = None) -> dict:
+    """
+    Get the training configuration from the model directory.
+    :param model_configuration: model configuration
+    :param training_id: training name
+    :return: training configuration
+    """
+    model_dir = model_configuration["model_dir"]
+    # get list of trainings
+    training_base_dir = os.path.join(model_dir, "trainings")
+    # if no training id is given, use the latest training
+    if training_id is None:
+        training_dirs = os.listdir(training_base_dir)
+        training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
+                         os.path.isdir(os.path.join(training_base_dir, dir))]
+        training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
+        training_dir = training_dirs[0]
+    else:
+        training_dir = os.path.join(training_base_dir, training_id)
+
+    training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
+    return training_configuration
+
+
+def find_models(base_dir: str, current_dir: str = None) -> dict:
+    """
+    Find all models in the base directory.
+    :param base_dir: base directory
+    :return: list of model directories
+    """
+    model_dirs = dict()
+    is_root = base_dir == current_dir
+    if current_dir is None:
+        current_dir = base_dir
+    for sub_dir_name in os.listdir(current_dir):
+        sub_dir = os.path.join(current_dir, sub_dir_name)
+        if not os.path.isdir(sub_dir):
+            continue
+        model_config_path = os.path.join(sub_dir, "model_configuration.pickle")
+
+        if current_dir not in model_dirs:
+            model_dirs[current_dir] = list()
+        if os.path.exists(model_config_path) and is_root:
+            model_dirs[current_dir].append(sub_dir)
+        elif os.path.exists(model_config_path) and not is_root:
+            model_dirs[current_dir].append(sub_dir)
+        else:
+            model_dirs = model_dirs | find_models(base_dir, sub_dir)
+
+    return model_dirs
+
+
+def apply_sigmoid_if_necessary(model_outputs: np.ndarray | torch.Tensor,
+                               model_configuration: dict) -> np.ndarray:
+    """
+    Applies a sigmoid to the outputs of a model, if they use a BCE loss.
+    Args:
+        model_outputs: models outputs
+        model_configuration: configuration of the model that produced the outputs
+
+    Returns:
+        processed_model_outputs: processed model outputs with sigmoid applied where necessary
+    """
+
+    target_features = model_configuration["feature_config"]["target_features"]
+    ignored_features = model_configuration["feature_config"]["ignored_features"]
+    used_targets = [feature for feature in target_features if feature["name"] not in ignored_features]
+
+    is_3d = len(model_outputs.shape) == 3
+
+    processed_outputs = np.empty_like(model_outputs)
+    for i, used_target in enumerate(used_targets):
+        if used_target["loss_fn"] == nn.BCEWithLogitsLoss:
+            if is_3d:
+                target_values = model_outputs[:, 0, i]
+                if isinstance(target_values, np.ndarray):
+                    target_values = torch.from_numpy(target_values)
+                processed_outputs[:, 0, i] = torch.sigmoid(target_values).numpy()
+            else:
+                target_values = model_outputs[:, i]
+                if isinstance(target_values, np.ndarray):
+                    target_values = torch.from_numpy(target_values)
+                processed_outputs[:, i] = torch.sigmoid(target_values).numpy()
+        else:
+            if is_3d:
+                processed_outputs[:, 0, i] = model_outputs[:, 0, i].numpy() if isinstance(model_outputs,
+                                                                                          torch.Tensor) else model_outputs[
+                                                                                                             :, 0, i]
+            else:
+                processed_outputs[:, i] = model_outputs[:, i].numpy() if isinstance(model_outputs,
+                                                                                    torch.Tensor) else model_outputs[:,
+                                                                                                       i]
+
+    return processed_outputs
+
+
+def scale_features(features: np.ndarray | torch.Tensor,
+                   feature_names: list | None,
+                   model_configuration: dict,
+                   scalers: dict = None,
+                   subset_name: str = "train") -> np.ndarray:
+    """
+    Scales the given features based on the scalers associated with the given model configuration
+    Args:
+        features: feature arrays, may be up to 3D (but may only have 2 usable dimensions, (B, 1, F) is allowed)
+        feature_names: names of the features in the same order as in the input array, may be None, then no scaling is performed
+        model_configuration: model configuration of the model the features belong to
+        scalers: dict of scales associated by name, will be fetched from model configuration, if None
+        subset_name: subset name, either "train", "test" or "val", used for determining the correct scalers
+
+    Returns:
+        features: scaled features
+    """
+
+    if features.shape[-1] != len(feature_names):
+        raise ValueError("Number of features does not match number of given feature names")
+
+    # remove singleton dimension, if 3D
+    if len(features.shape) == 3:
+        if features.shape[1] != 1:
+            raise ValueError(f"Input has 3 dimensions, expected (B, 1, F), but got {features.shape}")
+        features = features[:, 0, :]
+
+    if scalers is None:
+        scalers = get_scalers_for_model(model_configuration)
+
+    if isinstance(features, torch.Tensor):
+        features = features.detach().cpu().numpy()
+
+    scaled_features = np.empty_like(features)
+    for i in range(len(feature_names)):
+        # skip scaling, if feature name is None
+        if feature_names[i] is not None:
+            scaled_output = inverse_scale_feature(features[:, i],
+                                                  feature_names[i],
+                                                  scalers,
+                                                  subset_name=subset_name)
+        else:
+            scaled_output = features[:, i]
+
+        # make sure to make into 1D array when assigning
+        scaled_features[:, i] = scaled_output.ravel()
+
+    return scaled_features