Coverage for src / autoencodix / imagix.py: 93%

1from typing import Dict, Optional, Type, Union 

2from typing import Dict, Optional, Type, Union 

3import torch 

4import numpy as np 

5 

6from autoencodix.base._base_dataset import BaseDataset 

7from autoencodix.utils._utils import config_method 

8from autoencodix.base._base_loss import BaseLoss 

9from autoencodix.base._base_pipeline import BasePipeline 

10from autoencodix.base._base_trainer import BaseTrainer 

11from autoencodix.base._base_visualizer import BaseVisualizer 

12from autoencodix.base._base_preprocessor import BasePreprocessor 

13from autoencodix.base._base_autoencoder import BaseAutoencoder 

14from autoencodix.data._datasetcontainer import DatasetContainer 

15from autoencodix.data._datasplitter import DataSplitter 

16from autoencodix.data.datapackage import DataPackage 

17from autoencodix.data._image_dataset import ImageDataset 

18from autoencodix.data._image_processor import ImagePreprocessor 

19from autoencodix.evaluate._general_evaluator import GeneralEvaluator 

20from autoencodix.modeling._imagevae_architecture import ImageVAEArchitecture 

21from autoencodix.trainers._general_trainer import GeneralTrainer 

22from autoencodix.utils._result import Result 

23from autoencodix.configs.default_config import DefaultConfig 

24from autoencodix.utils._losses import VarixLoss 

25from autoencodix.visualize._imagix_visualizer import ImagixVisualizer 

26 

27 

28class Imagix(BasePipeline): 

29 """Imagix specific version of the BasePipeline class. 

30 

31 This class extends BasePipeline. See the parent class for a full list 

32 of attributes and methods. 

33 

34 Additional Attributes: 

35 _default_config: Is set to DefaultConfig here. 

36 

37 """ 

38 

39 def __init__( 

40 self, 

41 data: Optional[Union[DataPackage, DatasetContainer]] = None, 

42 trainer_type: Type[BaseTrainer] = GeneralTrainer, 

43 dataset_type: Type[BaseDataset] = ImageDataset, 

44 model_type: Type[BaseAutoencoder] = ImageVAEArchitecture, 

45 loss_type: Type[BaseLoss] = VarixLoss, 

46 preprocessor_type: Type[BasePreprocessor] = ImagePreprocessor, 

47 visualizer: Optional[Type[BaseVisualizer]] = ImagixVisualizer, 

48 evaluator: Optional[Type[GeneralEvaluator]] = GeneralEvaluator, 

49 result: Optional[Result] = None, 

50 datasplitter_type: Type[DataSplitter] = DataSplitter, 

51 custom_splits: Optional[Dict[str, np.ndarray]] = None, 

52 config: Optional[DefaultConfig] = None, 

53 ) -> None: 

54 """Initialize Imagix pipeline with customizable components. 

55 

56 Some components are passed as types rather than instances because they require 

57 data that is only available after preprocessing. 

58 

59 See Parentclass for full list of Args. 

60 

61 """ 

62 self._default_config = DefaultConfig() 

63 super().__init__( 

64 data=data, 

65 dataset_type=dataset_type, 

66 trainer_type=trainer_type, 

67 model_type=model_type, 

68 loss_type=loss_type, 

69 preprocessor_type=preprocessor_type, 

70 visualizer=visualizer, 

71 evaluator=evaluator, 

72 result=result, 

73 datasplitter_type=datasplitter_type, 

74 config=config, 

75 custom_split=custom_splits, 

76 )