from dataloaders.dataset1d import EcgDataset1D

from dataloaders.dataset2d import EcgDataset2D

from models import models1d, models2d

from runners.base_runner import BaseRunner

class Runner2D(BaseRunner):

    def __init__(self, config):

        super().__init__(config)

    def _init_net(self):

        model = getattr(models2d, self.config["model"])(

            num_classes=self.config["num_classes"],

        )

        model = model.to(self.config["device"])

        return model

    def _init_dataloader(self):

        inference_loader = EcgDataset2D(

            self.config["json"], self.config["mapping_json"],

        ).get_dataloader(

            batch_size=self.config["batch_size"],

            num_workers=self.config["num_workers"],

            shuffle=False,

        )

        return inference_loader

class Runner1D(BaseRunner):

    def __init__(self, config):

        super().__init__(config)

    def _init_net(self):

        model = getattr(models1d, self.config["model"])(

            num_classes=self.config["num_classes"],

        )

        model = model.to(self.config["device"])

        return model

    def _init_dataloader(self):

        inference_loader = EcgDataset1D(

            self.config["json"], self.config["mapping_json"],

        ).get_dataloader(

            batch_size=self.config["batch_size"],

            num_workers=self.config["num_workers"],

            shuffle=False,

        )

        return inference_loader