<reponame>Ramsha04/kits19-2d-reproduce<gh_stars>1-10
from glob import glob
from abc import abstractmethod
import os

import torch
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split

from kits19cnn.inference.utils import load_weights_infer
from kits19cnn.io import TestVoxelDataset
from catalyst.utils import get_device

class BaseInferenceExperiment(object):
    def __init__(self, config: dict):
        """
        Args:
            config (dict):

        Attributes:
            config-related:
                config (dict):
                io_params (dict):
                    in_dir (key: str): path to the data folder
                    test_size (key: float): split size for test
                    split_seed (key: int): seed
                    batch_size (key: int): <-
                    num_workers (key: int): # of workers for data loaders
            split_dict (dict): test_ids
            test_dset (torch.data.Dataset): <-
            loaders (dict): train/validation loaders
            model (torch.nn.Module): <-
        """
        # for reuse
        self.config = config
        self.io_params = config["io_params"]
        # initializing the experiment components
        self.case_list = self.setup_im_ids()
        test_ids = self.get_split()[-1] if config["with_masks"] else self.case_list
        print(f"Inferring on {len(test_ids)} test cases")
        self.test_dset = self.get_datasets(test_ids)
        self.loaders = self.get_loaders()
        self.model = self.get_model().to(get_device())
        self.load_weights()
        print(f"Device: {get_device()}")

    @abstractmethod
    def get_datasets(self, test_ids):
        """
        Initializes the data augmentation and preprocessing transforms. Creates
        and returns the train and validation datasets.
        """
        return

    @abstractmethod
    def get_model(self):
        """
        Creates and returns the model.
        """
        return

    def setup_im_ids(self):
        """
        Creates a list of all paths to case folders for the dataset split
        """
        search_path = os.path.join(self.config["in_dir"], "*/")
        case_list = sorted(glob(search_path))
        case_list = case_list[:210] if self.config["with_masks"] else case_list[210:]
        return case_list

    def get_split(self):
        """
        Creates train/valid filename splits
        """
        # setting up the train/val split with filenames
        split_seed: int = self.io_params["split_seed"]
        test_size: float = self.io_params["test_size"]
        if test_size == 1:
            return (None, None, self.case_list)
        else:
            # doing the splits: 1-test_size, test_size//2, test_size//2
            print("Splitting the dataset normally...")
            train_ids, total_test = train_test_split(self.case_list,
                                                     random_state=split_seed,
                                                     test_size=test_size)
            val_ids, test_ids = train_test_split(sorted(total_test),
                                                 random_state=split_seed,
                                                 test_size=0.5)
            return (train_ids, val_ids, test_ids)

    def get_loaders(self):
        """
        Creates train/val loaders from datasets created in self.get_datasets.
        Returns the loaders.
        """
        # setting up the loaders
        b_size, num_workers = self.io_params["batch_size"], self.io_params["num_workers"]
        test_loader = DataLoader(self.test_dset, batch_size=b_size,
                                  shuffle=False, num_workers=num_workers)
        return {"test": test_loader}

    def load_weights(self):
        """
        Loads model weights.
        """
        checkpoint_path = self.config["checkpoint_path"]
        print(f"Loading {checkpoint_path} into model...")
        self.model = load_weights_infer(checkpoint_path, self.model)

from glob import glob
from abc import abstractmethod
import os

import torch
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split

from kits19cnn.inference.utils import load_weights_infer
from kits19cnn.io import TestVoxelDataset
from catalyst.utils import get_device

class BaseInferenceExperiment(object):
    def __init__(self, config: dict):
        """
        Args:
            config (dict):

        Attributes:
            config-related:
                config (dict):
                io_params (dict):
                    in_dir (key: str): path to the data folder
                    test_size (key: float): split size for test
                    split_seed (key: int): seed
                    batch_size (key: int): <-
                    num_workers (key: int): # of workers for data loaders
            split_dict (dict): test_ids
            test_dset (torch.data.Dataset): <-
            loaders (dict): train/validation loaders
            model (torch.nn.Module): <-
        """
        # for reuse
        self.config = config
        self.io_params = config["io_params"]
        # initializing the experiment components
        self.case_list = self.setup_im_ids()
        test_ids = self.get_split()[-1] if config["with_masks"] else self.case_list
        print(f"Inferring on {len(test_ids)} test cases")
        self.test_dset = self.get_datasets(test_ids)
        self.loaders = self.get_loaders()
        self.model = self.get_model().to(get_device())
        self.load_weights()
        print(f"Device: {get_device()}")

    @abstractmethod
    def get_datasets(self, test_ids):
        """
        Initializes the data augmentation and preprocessing transforms. Creates
        and returns the train and validation datasets.
        """
        return

    @abstractmethod
    def get_model(self):
        """
        Creates and returns the model.
        """
        return

    def setup_im_ids(self):
        """
        Creates a list of all paths to case folders for the dataset split
        """
        search_path = os.path.join(self.config["in_dir"], "*/")
        case_list = sorted(glob(search_path))
        case_list = case_list[:210] if self.config["with_masks"] else case_list[210:]
        return case_list

    def get_split(self):
        """
        Creates train/valid filename splits
        """
        # setting up the train/val split with filenames
        split_seed: int = self.io_params["split_seed"]
        test_size: float = self.io_params["test_size"]
        if test_size == 1:
            return (None, None, self.case_list)
        else:
            # doing the splits: 1-test_size, test_size//2, test_size//2
            print("Splitting the dataset normally...")
            train_ids, total_test = train_test_split(self.case_list,
                                                     random_state=split_seed,
                                                     test_size=test_size)
            val_ids, test_ids = train_test_split(sorted(total_test),
                                                 random_state=split_seed,
                                                 test_size=0.5)
            return (train_ids, val_ids, test_ids)

    def get_loaders(self):
        """
        Creates train/val loaders from datasets created in self.get_datasets.
        Returns the loaders.
        """
        # setting up the loaders
        b_size, num_workers = self.io_params["batch_size"], self.io_params["num_workers"]
        test_loader = DataLoader(self.test_dset, batch_size=b_size,
                                  shuffle=False, num_workers=num_workers)
        return {"test": test_loader}

    def load_weights(self):
        """
        Loads model weights.
        """
        checkpoint_path = self.config["checkpoint_path"]
        print(f"Loading {checkpoint_path} into model...")
        self.model = load_weights_infer(checkpoint_path, self.model)
Args:             config (dict):          Attributes:             config-related:                 config (dict):                 io_params (dict):                     in_dir (key: str): path to the data folder                     test_size (key: float): split size for test                     split_seed (key: int): seed                     batch_size (key: int): <-                     num_workers (key: int): # of workers for data loaders             split_dict (dict): test_ids             test_dset (torch.data.Dataset): <-             loaders (dict): train/validation loaders             model (torch.nn.Module): <- # for reuse # initializing the experiment components Initializes the data augmentation and preprocessing transforms. Creates         and returns the train and validation datasets. Creates and returns the model. Creates a list of all paths to case folders for the dataset split Creates train/valid filename splits # setting up the train/val split with filenames # doing the splits: 1-test_size, test_size//2, test_size//2 Creates train/val loaders from datasets created in self.get_datasets.         Returns the loaders. # setting up the loaders Loads model weights.