dragon-ml-toolbox 19.14.0__py3-none-any.whl → 20.0.0__py3-none-any.whl

ml_tools/ML_models_vision/_image_classification.py

@@ -0,0 +1,182 @@
+from torch import nn
+import torchvision.models as vision_models
+from typing import Literal, Optional
+
+from ._base_wrapper import _BaseVisionWrapper
+
+
+__all__ = [
+    "DragonResNet",
+    "DragonEfficientNet",
+    "DragonVGG",
+]
+
+
+# Image classification
+class DragonResNet(_BaseVisionWrapper):
+    """
+    Image Classification
+    
+    A customizable wrapper for the torchvision ResNet family, compatible
+    with saving/loading architecture.
+
+    This wrapper allows for customizing the model backbone, input channels,
+    and the number of output classes for transfer learning.
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int = 3,
+                 model_name: Literal["resnet18", "resnet34", "resnet50", "resnet101", "resnet152"] = 'resnet50',
+                 init_with_pretrained: bool = False):
+        """
+        Args:
+            num_classes (int):
+                Number of output classes for the final layer.
+            in_channels (int):
+                Number of input channels (e.g., 1 for grayscale, 3 for RGB).
+            model_name (str):
+                The name of the ResNet model to use (e.g., 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'). Number is the layer count.
+            init_with_pretrained (bool):
+                If True, initializes the model with weights pretrained on ImageNet. This flag is for initialization only and is NOT saved in the architecture config.
+        """
+        
+        weights_enum_name = getattr(vision_models, f"{model_name.upper()}_Weights", None)
+        
+        super().__init__(
+            num_classes=num_classes,
+            in_channels=in_channels,
+            model_name=model_name,
+            init_with_pretrained=init_with_pretrained,
+            weights_enum_name=weights_enum_name
+        )
+
+    def _get_input_layer(self) -> nn.Conv2d:
+        return self.model.conv1
+
+    def _set_input_layer(self, layer: nn.Conv2d):
+        self.model.conv1 = layer
+
+    def _get_output_layer(self) -> Optional[nn.Linear]:
+        return self.model.fc
+
+    def _set_output_layer(self, layer: nn.Linear):
+        self.model.fc = layer
+
+
+class DragonEfficientNet(_BaseVisionWrapper):
+    """
+    Image Classification
+    
+    A customizable wrapper for the torchvision EfficientNet family, compatible
+    with saving/loading architecture.
+
+    This wrapper allows for customizing the model backbone, input channels,
+    and the number of output classes for transfer learning.
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int = 3,
+                 model_name: str = 'efficientnet_b0',
+                 init_with_pretrained: bool = False):
+        """
+        Args:
+            num_classes (int):
+                Number of output classes for the final layer.
+            in_channels (int):
+                Number of input channels (e.g., 1 for grayscale, 3 for RGB).
+            model_name (str):
+                The name of the EfficientNet model to use (e.g., 'efficientnet_b0'
+                through 'efficientnet_b7', or 'efficientnet_v2_s', 'efficientnet_v2_m', 'efficientnet_v2_l').
+            init_with_pretrained (bool):
+                If True, initializes the model with weights pretrained on
+                ImageNet. This flag is for initialization only and is
+                NOT saved in the architecture config. Defaults to False.
+        """
+        
+        weights_enum_name = getattr(vision_models, f"{model_name.upper()}_Weights", None)
+
+        super().__init__(
+            num_classes=num_classes,
+            in_channels=in_channels,
+            model_name=model_name,
+            init_with_pretrained=init_with_pretrained,
+            weights_enum_name=weights_enum_name
+        )
+
+    def _get_input_layer(self) -> nn.Conv2d:
+        # The first conv layer in EfficientNet is model.features[0][0]
+        return self.model.features[0][0]
+
+    def _set_input_layer(self, layer: nn.Conv2d):
+        self.model.features[0][0] = layer
+
+    def _get_output_layer(self) -> Optional[nn.Linear]:
+        # The classifier in EfficientNet is model.classifier[1]
+        if hasattr(self.model, 'classifier') and isinstance(self.model.classifier, nn.Sequential):
+            output_layer = self.model.classifier[1]
+            if isinstance(output_layer, nn.Linear):
+                return output_layer
+        return None
+
+    def _set_output_layer(self, layer: nn.Linear):
+        self.model.classifier[1] = layer
+
+
+class DragonVGG(_BaseVisionWrapper):
+    """
+    Image Classification
+    
+    A customizable wrapper for the torchvision VGG family, compatible
+    with saving/loading architecture.
+
+    This wrapper allows for customizing the model backbone, input channels,
+    and the number of output classes for transfer learning.
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int = 3,
+                 model_name: Literal["vgg11", "vgg13", "vgg16", "vgg19", "vgg11_bn", "vgg13_bn", "vgg16_bn", "vgg19_bn"] = 'vgg16',
+                 init_with_pretrained: bool = False):
+        """
+        Args:
+            num_classes (int):
+                Number of output classes for the final layer.
+            in_channels (int):
+                Number of input channels (e.g., 1 for grayscale, 3 for RGB).
+            model_name (str):
+                The name of the VGG model to use (e.g., 'vgg16', 'vgg16_bn').
+            init_with_pretrained (bool):
+                If True, initializes the model with weights pretrained on
+                ImageNet. This flag is for initialization only and is
+                NOT saved in the architecture config. Defaults to False.
+        """
+        
+        # Format model name to find weights enum, e.g., vgg16_bn -> VGG16_BN_Weights
+        weights_enum_name = f"{model_name.replace('_bn', '_BN').upper()}_Weights"
+        
+        super().__init__(
+            num_classes=num_classes,
+            in_channels=in_channels,
+            model_name=model_name,
+            init_with_pretrained=init_with_pretrained,
+            weights_enum_name=weights_enum_name
+        )
+
+    def _get_input_layer(self) -> nn.Conv2d:
+        # The first conv layer in VGG is model.features[0]
+        return self.model.features[0]
+
+    def _set_input_layer(self, layer: nn.Conv2d):
+        self.model.features[0] = layer
+
+    def _get_output_layer(self) -> Optional[nn.Linear]:
+        # The final classifier in VGG is model.classifier[6]
+        if hasattr(self.model, 'classifier') and isinstance(self.model.classifier, nn.Sequential) and len(self.model.classifier) == 7:
+            output_layer = self.model.classifier[6]
+            if isinstance(output_layer, nn.Linear):
+                return output_layer
+        return None
+
+    def _set_output_layer(self, layer: nn.Linear):
+        self.model.classifier[6] = layer
+

ml_tools/ML_models_vision/_image_segmentation.py

@@ -0,0 +1,108 @@
+from torch import nn
+from typing import Literal
+
+from ._base_wrapper import _BaseSegmentationWrapper
+
+
+__all__ = [
+    "DragonFCN",
+    "DragonDeepLabv3",
+]
+
+
+class DragonFCN(_BaseSegmentationWrapper):
+    """
+    Image Segmentation
+    
+    A customizable wrapper for the torchvision FCN (Fully Convolutional Network)
+    family, compatible with saving/loading architecture.
+
+    This wrapper allows for customizing the model backbone, input channels,
+    and the number of output classes for transfer learning.
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int = 3,
+                 model_name: Literal["fcn_resnet50", "fcn_resnet101"] = 'fcn_resnet50',
+                 init_with_pretrained: bool = False):
+        """
+        Args:
+            num_classes (int):
+                Number of output classes (including background).
+            in_channels (int):
+                Number of input channels (e.g., 1 for grayscale, 3 for RGB).
+            model_name (str):
+                The name of the FCN model to use ('fcn_resnet50' or 'fcn_resnet101').
+            init_with_pretrained (bool):
+                If True, initializes the model with weights pretrained on COCO.
+                This flag is for initialization only and is NOT saved in the
+                architecture config. Defaults to False.
+        """
+        # Format model name to find weights enum, e.g., fcn_resnet50 -> FCN_ResNet50_Weights
+        weights_model_name = model_name.replace('fcn_', 'FCN_').replace('resnet', 'ResNet')
+        weights_enum_name = f"{weights_model_name}_Weights"
+        
+        super().__init__(
+            num_classes=num_classes,
+            in_channels=in_channels,
+            model_name=model_name,
+            init_with_pretrained=init_with_pretrained,
+            weights_enum_name=weights_enum_name
+        )
+
+    def _get_input_layer(self) -> nn.Conv2d:
+        # FCN models use a ResNet backbone, input layer is backbone.conv1
+        return self.model.backbone.conv1
+
+    def _set_input_layer(self, layer: nn.Conv2d):
+        self.model.backbone.conv1 = layer
+
+
+class DragonDeepLabv3(_BaseSegmentationWrapper):
+    """
+    Image Segmentation
+    
+    A customizable wrapper for the torchvision DeepLabv3 family, compatible
+    with saving/loading architecture.
+
+    This wrapper allows for customizing the model backbone, input channels,
+    and the number of output classes for transfer learning.
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int = 3,
+                 model_name: Literal["deeplabv3_resnet50", "deeplabv3_resnet101"] = 'deeplabv3_resnet50',
+                 init_with_pretrained: bool = False):
+        """
+        Args:
+            num_classes (int):
+                Number of output classes (including background).
+            in_channels (int):
+                Number of input channels (e.g., 1 for grayscale, 3 for RGB).
+            model_name (str):
+                The name of the DeepLabv3 model to use ('deeplabv3_resnet50' or 'deeplabv3_resnet101').
+            init_with_pretrained (bool):
+                If True, initializes the model with weights pretrained on COCO.
+                This flag is for initialization only and is NOT saved in the
+                architecture config. Defaults to False.
+        """
+        
+        # Format model name to find weights enum, e.g., deeplabv3_resnet50 -> DeepLabV3_ResNet50_Weights
+        weights_model_name = model_name.replace('deeplabv3_', 'DeepLabV3_').replace('resnet', 'ResNet')
+        weights_enum_name = f"{weights_model_name}_Weights"
+        
+        super().__init__(
+            num_classes=num_classes,
+            in_channels=in_channels,
+            model_name=model_name,
+            init_with_pretrained=init_with_pretrained,
+            weights_enum_name=weights_enum_name
+        )
+
+    def _get_input_layer(self) -> nn.Conv2d:
+        # DeepLabv3 models use a ResNet backbone, input layer is backbone.conv1
+        return self.model.backbone.conv1
+
+    def _set_input_layer(self, layer: nn.Conv2d):
+        self.model.backbone.conv1 = layer
+

ml_tools/ML_models_vision/_imprimir.py

@@ -0,0 +1,16 @@
+from .._core import _imprimir_disponibles
+
+_GRUPOS = [
+    # Image Classification
+    "DragonResNet",
+    "DragonEfficientNet",
+    "DragonVGG",
+    # Image Segmentation
+    "DragonFCN",
+    "DragonDeepLabv3",
+    # Object Detection
+    "DragonFastRCNN",
+]
+
+def info():
+    _imprimir_disponibles(_GRUPOS)

ml_tools/ML_models_vision/_object_detection.py

@@ -0,0 +1,135 @@
+import torch
+from torch import nn
+from torchvision.models import detection as detection_models
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from typing import Any, Literal, Optional
+
+from ..ML_models._base_save_load import _ArchitectureHandlerMixin
+
+from .._core import get_logger
+
+
+_LOGGER = get_logger("DragonFastRCNN")
+
+
+__all__ = [
+    "DragonFastRCNN",
+]
+
+
+# Object Detection
+class DragonFastRCNN(nn.Module, _ArchitectureHandlerMixin):
+    """
+    Object Detection
+    
+    A customizable wrapper for the torchvision Faster R-CNN family.
+    
+    This wrapper allows for customizing the model backbone, input channels,
+    and the number of output classes for transfer learning.
+
+    NOTE: Use an Object Detection compatible trainer.
+    """
+    def __init__(self,
+                 num_classes: int,
+                 in_channels: int = 3,
+                 model_name: Literal["fasterrcnn_resnet50_fpn", "fasterrcnn_resnet50_fpn_v2"] = 'fasterrcnn_resnet50_fpn_v2',
+                 init_with_pretrained: bool = False):
+        """
+        Args:
+            num_classes (int):
+                Number of output classes (including background).
+            in_channels (int):
+                Number of input channels (e.g., 1 for grayscale, 3 for RGB).
+            model_name (str):
+                The name of the Faster R-CNN model to use.
+            init_with_pretrained (bool):
+                If True, initializes the model with weights pretrained on COCO.
+                This flag is for initialization only and is NOT saved in the
+                architecture config. Defaults to False.
+        """
+        super().__init__()
+        
+        # --- 1. Validation and Configuration ---
+        if not hasattr(detection_models, model_name):
+            _LOGGER.error(f"'{model_name}' is not a valid model name in torchvision.models.detection.")
+            raise ValueError()
+
+        self.num_classes = num_classes
+        self.in_channels = in_channels
+        self.model_name = model_name
+        self._pretrained_default_transforms = None
+
+        # --- 2. Instantiate the base model ---
+        model_constructor = getattr(detection_models, model_name)
+        
+        # Format model name to find weights enum, e.g., fasterrcnn_resnet50_fpn_v2 -> FasterRCNN_ResNet50_FPN_V2_Weights
+        weights_model_name = model_name.replace('fasterrcnn_', 'FasterRCNN_').replace('resnet', 'ResNet').replace('_fpn', '_FPN')
+        weights_enum_name = f"{weights_model_name.upper()}_Weights"
+        
+        weights_enum = getattr(detection_models, weights_enum_name, None) if weights_enum_name else None
+        weights = weights_enum.DEFAULT if weights_enum and init_with_pretrained else None
+        
+        if weights:
+            self._pretrained_default_transforms = weights.transforms()
+
+        self.model = model_constructor(weights=weights, weights_backbone=weights)
+        
+        # --- 4. Modify the output layer (Box Predictor) ---
+        # Get the number of input features for the classifier
+        in_features = self.model.roi_heads.box_predictor.cls_score.in_features
+        # Replace the pre-trained head with a new one
+        self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+
+        # --- 3. Modify the input layer (Backbone conv1) ---
+        if in_channels != 3:
+            original_conv1 = self.model.backbone.body.conv1
+            
+            new_conv1 = nn.Conv2d(
+                in_channels,
+                original_conv1.out_channels,
+                kernel_size=original_conv1.kernel_size, # type: ignore
+                stride=original_conv1.stride, # type: ignore
+                padding=original_conv1.padding, # type: ignore
+                bias=(original_conv1.bias is not None)
+            )
+            
+            # (Optional) Average original weights if starting from pretrained
+            if init_with_pretrained and original_conv1.in_channels == 3 and weights is not None:
+                with torch.no_grad():
+                    # Average the weights across the input channel dimension
+                    avg_weights = torch.mean(original_conv1.weight, dim=1, keepdim=True)
+                    # Repeat the averaged weights for the new number of input channels
+                    new_conv1.weight[:] = avg_weights.repeat(1, in_channels, 1, 1)
+
+            self.model.backbone.body.conv1 = new_conv1
+
+    def forward(self, images: list[torch.Tensor], targets: Optional[list[dict[str, torch.Tensor]]] = None):
+        """
+        Defines the forward pass.
+        
+        - In train mode, expects (images, targets) and returns a dict of losses.
+        - In eval mode, expects (images) and returns a list of prediction dicts.
+        """
+        # The model's forward pass handles train/eval mode internally.
+        return self.model(images, targets)
+
+    def get_architecture_config(self) -> dict[str, Any]:
+        """
+        Returns the structural configuration of the model.
+        The 'init_with_pretrained' flag is intentionally omitted,
+        as .load() should restore the architecture, not the weights.
+        """
+        return {
+            'num_classes': self.num_classes,
+            'in_channels': self.in_channels,
+            'model_name': self.model_name
+        }
+
+    def __repr__(self) -> str:
+        """Returns the developer-friendly string representation of the model."""
+        return (
+            f"{self.__class__.__name__}(model='{self.model_name}', "
+            f"in_channels={self.in_channels}, "
+            f"num_classes={self.num_classes})"
+        )
+

ml_tools/ML_optimization/__init__.py

@@ -0,0 +1,21 @@
+from ._multi_dragon import DragonParetoOptimizer
+
+from ._single_dragon import DragonOptimizer
+
+from ._single_manual import (
+    FitnessEvaluator,
+    create_pytorch_problem,
+    run_optimization,
+)
+
+from ._imprimir import info
+
+
+__all__ = [
+    "DragonParetoOptimizer",
+    "DragonOptimizer",
+    # manual optimization tools
+    "FitnessEvaluator",
+    "create_pytorch_problem",
+    "run_optimization",
+]

ml_tools/ML_optimization/_imprimir.py

@@ -0,0 +1,13 @@
+from .._core import _imprimir_disponibles
+
+_GRUPOS = [
+    "DragonParetoOptimizer",
+    "DragonOptimizer",
+    # manual optimization tools
+    "FitnessEvaluator",
+    "create_pytorch_problem",
+    "run_optimization",
+]
+
+def info():
+    _imprimir_disponibles(_GRUPOS)

ml_tools/{_core/_ML_optimization_pareto.py → ML_optimization/_multi_dragon.py}

@@ -1,13 +1,10 @@
 import torch
-import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
-import matplotlib.colors as mcolors
-import matplotlib.cm as cm
 from matplotlib.collections import LineCollection
 import seaborn as sns
 from pathlib import Path
-from typing import Literal, Union, List, Optional, Dict
+from typing import Literal, Union, Optional
 from tqdm import tqdm
 import plotly.express as px
 import plotly.graph_objects as go
@@ -17,22 +14,22 @@ from evotorch import Problem
 from evotorch.operators import SimulatedBinaryCrossOver, GaussianMutation
 from evotorch.operators import functional as func_ops
 
-from ._SQL import DragonSQL
-from ._ML_inference import DragonInferenceHandler
-from ._ML_chaining_inference import DragonChainInference
-from ._ML_configuration import DragonParetoConfig
-from ._optimization_tools import create_optimization_bounds, plot_optimal_feature_distributions_from_dataframe, load_continuous_bounds_template
-from ._math_utilities import discretize_categorical_values
-from ._utilities import save_dataframe_filename
-from ._IO_tools import save_json
-from ._path_manager import make_fullpath, sanitize_filename
-from ._logger import get_logger
-from ._script_info import _script_info
-from ._keys import PyTorchInferenceKeys, MLTaskKeys, ParetoOptimizationKeys
-from ._schema import FeatureSchema
+from ..SQL import DragonSQL
+from ..ML_inference import DragonInferenceHandler
+from ..ML_inference._chain_inference import DragonChainInference
+from ..ML_configuration import DragonParetoConfig
+from ..optimization_tools import create_optimization_bounds, plot_optimal_feature_distributions_from_dataframe, load_continuous_bounds_template
+from ..utilities import save_dataframe_filename
+from ..IO_tools import save_json
+from ..schema import FeatureSchema
 
+from ..math_utilities import discretize_categorical_values
+from ..path_manager import make_fullpath, sanitize_filename
+from .._core import get_logger
+from ..keys._keys import PyTorchInferenceKeys, MLTaskKeys, ParetoOptimizationKeys
 
-_LOGGER = get_logger("Pareto Optimizer")
+
+_LOGGER = get_logger("DragonParetoOptimizer")
 
 
 __all__ = [
@@ -793,9 +790,9 @@ class _ParetoFitnessEvaluator:
     """
     def __init__(self,
                  inference_handler: Union[DragonInferenceHandler, DragonChainInference],
-                 target_indices: List[int],
-                 target_names: List[str],
-                 categorical_index_map: Optional[Dict[int, int]] = None,
+                 target_indices: list[int],
+                 target_names: list[str],
+                 categorical_index_map: Optional[dict[int, int]] = None,
                  discretize_start_at_zero: bool = True,
                  is_chain: bool = False):
         
@@ -856,6 +853,3 @@ class _ParetoFitnessEvaluator:
             # Slice specific indices -> (Batch, N_Selected_Targets)
             return preds[:, self.target_indices]
 
-
-def info():
-    _script_info(__all__)