openmlcore 1.0.0__tar.gz

openmlcore-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Votre Nom
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

openmlcore-1.0.0/PKG-INFO

@@ -0,0 +1,76 @@
+Metadata-Version: 2.4
+Name: openmlcore
+Version: 1.0.0
+Summary: OpenMLCore is an open source framework (wrapper which simplify the use of PyTorch)
+Author-email: Doyez Sébastien <doyez.sebastien34090@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 Votre Nom
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+Project-URL: Homepage, https://sebastien-doyez2812.github.io/openMLCore.github.io/
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10.0
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: tqdm
+Requires-Dist: onnx
+Requires-Dist: onnxscript
+Requires-Dist: torchinfo
+Requires-Dist: datasets
+Dynamic: license-file
+
+# OpenML-core, Open source Package for Machine learning
+
+
+## Why OpenML?
+Today's computer vision ecosystem is heavily fragmented. If you want YOLO, Ultralytics makes it trivial. But the moment you need advanced semantic segmentation, hybrid transformer backbones, or efficient fine-tuning on modern architectures, you end up stitching together disjointed codebases, mismatched tensor shapes, and unmaintained repositories.
+
+OpenML bridges this gap by bringing a clean, Hugging Face-inspired Developer Experience (DX) to advanced computer vision tasks—from classical models to SOTA transformers and parameter-efficient fine-tuning.
+## Supported Architectures:
+
+- Classic & Deep Segmentation: U-Net, Attention U-Net, UNet++, Mask R-CNN
+
+- Advanced Transformers: Swin Transformer, TransUNet, MobileUNetR
+
+- SOTA & Edge: Vision Transformers, SAM 2 adaptations via efficient Low-Rank Adaptation (LoRA)
+## Key features:
+
+Unified Object-Oriented API: Every model inherits from a clean ModelBase class supporting standard .train(), .eval(), and .help() workflows.
+
+Native LoRA Wrappers: Fine-tune massive vision models and transformers on modest hardware (like a single RTX 3070) without blowing up your VRAM.
+
+Hugging Face Hub Integration: Push and pull checkpoints effortlessly using native .safetensors formatting.
+
+MLflow Tracking Out-of-the-Box: Keep tabs on your losses, metrics, and hyperparameter logs without writing boilerplate monitoring code.
+
+Interactive Help System: Running model.help() directly opens up the exact architectural documentation page on the web.
+
+## Contributions:
+
+Contributions, feature requests, and bug reports are welcome! Feel free to open an Issue or submit a Pull Request on GitHub.
+
+## Author
+
+Sébastien Doyez
+ML & AI Engineer,
+Interested in Agentics, Computer Vision, Robotics & Deep Learning

openmlcore-1.0.0/README.md

@@ -0,0 +1,35 @@
+# OpenML-core, Open source Package for Machine learning
+
+
+## Why OpenML?
+Today's computer vision ecosystem is heavily fragmented. If you want YOLO, Ultralytics makes it trivial. But the moment you need advanced semantic segmentation, hybrid transformer backbones, or efficient fine-tuning on modern architectures, you end up stitching together disjointed codebases, mismatched tensor shapes, and unmaintained repositories.
+
+OpenML bridges this gap by bringing a clean, Hugging Face-inspired Developer Experience (DX) to advanced computer vision tasks—from classical models to SOTA transformers and parameter-efficient fine-tuning.
+## Supported Architectures:
+
+- Classic & Deep Segmentation: U-Net, Attention U-Net, UNet++, Mask R-CNN
+
+- Advanced Transformers: Swin Transformer, TransUNet, MobileUNetR
+
+- SOTA & Edge: Vision Transformers, SAM 2 adaptations via efficient Low-Rank Adaptation (LoRA)
+## Key features:
+
+Unified Object-Oriented API: Every model inherits from a clean ModelBase class supporting standard .train(), .eval(), and .help() workflows.
+
+Native LoRA Wrappers: Fine-tune massive vision models and transformers on modest hardware (like a single RTX 3070) without blowing up your VRAM.
+
+Hugging Face Hub Integration: Push and pull checkpoints effortlessly using native .safetensors formatting.
+
+MLflow Tracking Out-of-the-Box: Keep tabs on your losses, metrics, and hyperparameter logs without writing boilerplate monitoring code.
+
+Interactive Help System: Running model.help() directly opens up the exact architectural documentation page on the web.
+
+## Contributions:
+
+Contributions, feature requests, and bug reports are welcome! Feel free to open an Issue or submit a Pull Request on GitHub.
+
+## Author
+
+Sébastien Doyez
+ML & AI Engineer,
+Interested in Agentics, Computer Vision, Robotics & Deep Learning

openmlcore-1.0.0/pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openmlcore"               
+version = "1.0.0"                 
+authors = [
+  { name="Doyez Sébastien", email="doyez.sebastien34090@gmail.com" },
+]
+description = "OpenMLCore is an open source framework (wrapper which simplify the use of PyTorch)"
+readme = "README.md"
+requires-python = ">=3.10.0"
+license = { file = "LICENSE" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+"torch",
+"tqdm",
+"onnx",
+"onnxscript",
+"torchinfo",
+"datasets"
+]
+
+[project.urls]
+"Homepage" = "https://sebastien-doyez2812.github.io/openMLCore.github.io/"
+
+[tool.setuptools.packages.find]
+where = ["src"]

openmlcore-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

openmlcore-1.0.0/src/openmlcore/base/data.py

@@ -0,0 +1,48 @@
+# author: Sebastien Doyez
+# data is the file where we create custom Dataset class
+# SegmentationDataset is the dataset class for 1 hot encoding dataset
+
+import os
+from PIL import Image
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import v2
+from torchvision import tv_tensors
+
+class SegmentationDataset(Dataset):
+    def __init__(self, root_dir="data/train", transform=None, num_classes = 8):
+        self.root_dir = root_dir
+        self.images_dir = os.path.join(root_dir, "imgs")
+        self.masks_dir = os.path.join(root_dir, "labels")
+        self.transform = transform
+        self.num_classes = num_classes
+
+        self.img_names = sorted(os.listdir(self.images_dir))
+        self.mask_names = sorted(os.listdir(self.masks_dir))
+
+    def __len__(self):
+        return len(self.img_names)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.images_dir, self.img_names[idx])
+        mask_path = os.path.join(self.masks_dir, self.mask_names[idx])
+
+        image = Image.open(img_path).convert("RGB")
+        mask = Image.open(mask_path)
+
+        mask_np = np.array(mask, dtype = np.uint8)
+        masks = [((mask_np >> i ) & 1) for i in range (self.num_classes)]
+
+        mask_stacked = np.stack(masks, axis = 0).astype(np.float32)
+        mask_tensor = tv_tensors.Mask(mask_stacked)
+
+        if self.transform is not None:
+            image, mask_tensor = self.transform(image, mask_tensor)
+        else:
+            image = v2.functional.to_image(image)
+            image = v2.functional.to_dtype(image, dtype=torch.float32, scale=True)
+
+        return image, mask_tensor
+
+

openmlcore-1.0.0/src/openmlcore/loss/loss.py

@@ -0,0 +1,210 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class BaseLoss(nn.Module):
+    def __init__(self, name: str = "BaseLoss", doc_url: str = ""):
+        super().__init__()
+        self.name = name
+        self.smooth = 1e-6
+        self.doc_url = doc_url
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    def documentation(self) -> str:
+        print(f"Documentation for {self.name} is available here:\n{self.doc_url}")
+        try:
+            import webbrowser
+            webbrowser.open(self.doc_url)
+        except Exception as e:
+            print(f"[-] {e}")
+
+    
+#############################
+#   Class Loss Functions    #
+#############################
+
+class BCEWithLogitsLoss(BaseLoss):
+    def __init__(self, weights=None):
+        super().__init__(name="BCEWithLogitsLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        pos_weights = None
+        if weights is not None:
+            if not isinstance(weights, torch.Tensor):
+                pos_weights = torch.tensor(weights, device=device, dtype=torch.float32)
+            else:
+                pos_weights = weights.to(device)
+
+            if pos_weights.ndim == 1:
+                pos_weights = pos_weights.view(-1, 1, 1)
+        self.loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weights)
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        return self.loss_fn(outputs, targets.float())
+
+class DiceLoss(BaseLoss):
+    def __init__(self, weights = None):
+        super().__init__(name="DiceLoss", doc_url="https://arxiv.org/abs/1606.04797")
+        self.weights = weights
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        probs = torch.sigmoid(outputs)
+        probs   = probs.view  (probs.size(0)  , probs.size(1), -1)
+        targets = targets.view(targets.size(0), targets.size(1), -1)
+        
+        intersection = (probs * targets).sum(dim=2)
+        union        = probs.sum(dim=2) + targets.sum(dim=2) 
+
+        dice_score = (2.0 * intersection + self.smooth) / (union + self.smooth)
+        dice_loss = 1.0 - dice_score
+
+        if self.weights is not None:
+            if not isinstance(self.weights, torch.Tensor):
+                weights_tensor = torch.tensor(self.weights, device=outputs.device, dtype=outputs.dtype)
+            else:
+                weights_tensor = self.weights.to(outputs.device)
+
+            dice_loss = dice_loss * weights_tensor.unsqueeze(0)
+
+            return dice_loss.sum(dim=1).mean()
+        return dice_loss.mean()
+    
+# TODO:
+# 1 Hot encoding loss
+# 
+
+
+class MeanSquaredErrorLoss(BaseLoss):
+    def __init__(self, index_ignore=255):
+        super().__init__(name="MeanSquaredErrorLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html")
+        self.index_ignore = index_ignore
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        num_classes = outputs.shape[1]
+        probs = F.softmax(outputs, dim=1)
+        targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
+        
+        mse = (probs - targets_one_hot) ** 2
+        mse = mse * valid_mask
+        
+        num_valids = valid_mask.sum() * num_classes
+        return mse.sum() / torch.clamp(num_valids, min=1.0)
+
+
+
+class FocalLoss(BaseLoss):
+    def __init__(self, alpha=0.25, gamma=2.0, index_ignore=255):
+        super().__init__(name="FocalLoss", doc_url="https://arxiv.org/abs/1708.02002")
+        self.alpha = alpha
+        self.gamma = gamma
+        self.index_ignore = index_ignore
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        num_classes = outputs.shape[1]
+        targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
+        
+        bce_loss = F.binary_cross_entropy_with_logits(outputs, targets_one_hot, reduction='none')
+        pt = torch.exp(-bce_loss)
+        focal_loss = self.alpha * ((1 - pt) ** self.gamma) * bce_loss
+        focal_loss = focal_loss * valid_mask
+        num_valids = valid_mask.sum() * num_classes
+        return focal_loss.sum() / torch.clamp(num_valids, min=1.0)
+    
+class MeanSquaredErrorLoss(BaseLoss):
+    def __init__(self, index_ignore=255):
+        super().__init__(name="MeanSquaredErrorLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html")
+        self.index_ignore = index_ignore
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        num_classes = outputs.shape[1]
+        probs = F.softmax(outputs, dim=1)
+        targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
+        
+        mse = (probs - targets_one_hot) ** 2
+        mse = mse * valid_mask
+        
+        num_valids = valid_mask.sum() * num_classes
+        return mse.sum() / torch.clamp(num_valids, min=1.0)
+
+class HuberLoss(BaseLoss):
+    def __init__(self, delta=1.0, index_ignore=255):
+        super().__init__(name="HuberLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.HuberLoss.html")
+        self.delta = delta
+        self.index_ignore = index_ignore
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        num_classes = outputs.shape[1]
+        probs = F.softmax(outputs, dim=1)
+        targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
+        
+        error = probs - targets_one_hot
+        abs_error = torch.abs(error)
+        
+        huber = torch.where(
+            abs_error < self.delta,
+            0.5 * (error ** 2),
+            self.delta * (abs_error - 0.5 * self.delta)
+        )
+        huber = huber * valid_mask
+        
+        num_valids = valid_mask.sum() * num_classes
+        return huber.sum() / torch.clamp(num_valids, min=1.0)
+    
+class KLDivergenceLoss(BaseLoss):
+    def __init__(self, index_ignore=255):
+        super().__init__(name="KLDivergenceLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html")
+        self.eps = 1e-8
+        self.index_ignore = index_ignore
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        num_classes = outputs.shape[1]
+        log_probs = F.log_softmax(outputs, dim=1)
+        targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
+        
+        # KL-Div = p(x) * (log(p(x)) - log(q(x)))
+        # Pour une distribution One-Hot, cela se réduit à : - log(q(x)) sur la classe vraie
+        kl = F.kl_div(log_probs, targets_one_hot, reduction='none')
+        kl = kl * valid_mask
+        
+        num_valids = valid_mask.sum() * num_classes
+        return kl.sum() / torch.clamp(num_valids, min=1.0)
+
+####################################
+#   Custom Loss implementation     #
+####################################
+
+class LossFactory:
+    @staticmethod
+    def create(name: str, **kwargs ) -> BaseLoss:
+        match name:
+            case "BCE":
+                return BCEWithLogitsLoss(**kwargs)
+            case "MeanSquaredErrorLoss":
+                return MeanSquaredErrorLoss(**kwargs)
+            case "DiceLoss":
+                return DiceLoss(**kwargs)
+            case "FocalLoss":
+                return FocalLoss(**kwargs)
+            case "HuberLoss":
+                return HuberLoss(**kwargs)
+            case "KLDivergenceLoss":
+                return KLDivergenceLoss(**kwargs)
+            case _:
+                raise ValueError(f"Unknown loss function: {name}")
+
+class CustomLoss(BaseLoss):
+    def __init__(self, loss_fcns: list[BaseLoss], name: str, coefficients: list[float]):
+        assert len(loss_fcns) == len(coefficients), "The number of loss functions must match the number of coefficients."
+        super().__init__(name=name, doc_url=None)
+        self.loss_fcns = nn.ModuleList(loss_fcns)
+        self.coefficients = coefficients
+
+    def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        total_loss = 0.0
+        for loss_fn, coef in zip(self.loss_fcns, self.coefficients):
+            total_loss += coef * loss_fn(outputs, targets)
+        return total_loss
+    

openmlcore-1.0.0/src/openmlcore/metrics/metrics.py

@@ -0,0 +1,81 @@
+###             GUIDELINES            ###
+
+# The metrics are define by functions.
+# The function should ALWAYS return a list with the value of the metric of each class.
+# For example IoU (Intersection of Union) return a list.
+# Let's say we have 3 classes:
+# Class 0: sky
+# Class 1: Road
+# Class 2: Car
+# We should return a list of 3 elements like [IoU_sky, IoU_Road, IoU_Car]
+import torch
+
+def accuracy(y_pred, y_true):
+    num_batches = y_true.shape[0]
+    num_classes = y_true.shape[1]
+
+    metric_to_return = [0.0] * num_classes
+    for idx_current_batch in range(num_batches):
+        current_metrics = []
+        for idx_class in range(num_classes):
+            pred = (y_pred[idx_current_batch, idx_class, :, :] > 0.0).float()
+            gt   = y_true[idx_current_batch, idx_class, :, :]
+
+            correct = (pred == gt).float()
+            accuracy = (pred == gt).float().mean().item()
+
+            current_metrics.append(accuracy)
+        
+        for i in range(num_classes):
+            metric_to_return[i] += current_metrics[i]
+
+    final_metric = [0.0] * num_classes
+    for idx, metric in enumerate(metric_to_return):
+        final_metric[idx] = metric / num_batches 
+    return final_metric
+    
+def iou(y_pred, y_true,smooth = 1e-6):
+    preds = (y_pred > 0.0).float()
+    preds = preds.view(preds.size(0), preds.size(1), -1)
+    y_true = y_true.view(y_true.size(0), y_true.size(1), -1)
+
+    intersection = (preds * y_true).sum(dim=2)
+    total = preds.sum(dim=2) + y_true.sum(dim=2)
+    union = total - intersection
+
+    iou_per_class = (intersection + smooth) / (union + smooth)
+    mean_iou_per_class = iou_per_class.mean(dim=0)
+    
+    return mean_iou_per_class.tolist()
+
+
+def F1(y_pred, y_true, eps = 1e-7):
+    # F1 = 2. precision.recall/(precision + recall)
+    recall_res    = recall(y_pred=y_pred, y_true=y_true)
+    precision_res = precision(y_pred=y_pred, y_true=y_true)
+
+    F1_values = []
+    for idx in range(len(recall_res)):
+        current_F1 = 2 * recall_res[idx] * precision_res[idx] / (precision_res[idx] + recall_res[idx] + eps)
+        F1_values.append(current_F1)
+    return F1_values
+
+def recall(y_pred, y_true, threshold = 0.5, eps = 1e-7):
+    # Recall = TP/(TP+FN)
+    y_pred_bin = (y_pred >threshold).float()
+    y_true     = y_true.float()
+    # (0, 2, 3) means we are going to add the batch, H and W, we only keep the class
+    TP = torch.sum(y_pred_bin * y_true      , dim= (0, 2, 3)) 
+    FN = torch.sum(y_true * (1 - y_pred_bin), dim= (0, 2, 3))
+    precision = TP/(TP + FN + eps)
+    return precision.tolist()
+
+def precision(y_pred, y_true, threshold = 0.5, eps = 1e-7):
+    # Prec = TP/(TP+FP)
+    y_pred_bin = (y_pred >threshold).float()
+    y_true     = y_true.float()
+    # (0, 2, 3) means we are going to add the batch, H and W, we only keep the class
+    TP = torch.sum(y_pred_bin * y_true      , dim= (0, 2, 3)) 
+    FP = torch.sum((1 - y_true) * y_pred_bin, dim= (0, 2, 3))
+    precision = TP/(TP + FP + eps)
+    return precision.tolist()

openmlcore-1.0.0/src/openmlcore/models/UNets.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+from .base_model import BaseModel
+from .lib.convolutions import DoubleConv
+
+class UnetTorch(nn.Module):
+    def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64,
+                 kernel_size=3, padding=1):
+        super(UnetTorch, self).__init__()
+        self.downs = nn.ModuleList()
+        self.ups = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        in_c = input_channels
+        out_c = initial_filters
+        for _ in range(depth):
+            self.downs.append(DoubleConv(in_c, out_c, kernel_size=kernel_size, padding=padding))
+            in_c = out_c
+            out_c *= 2
+
+        self.bottleneck = DoubleConv(in_c, out_c, kernel_size=kernel_size, padding=padding)
+
+        for _ in range(depth):
+            self.ups.append(nn.ConvTranspose2d(out_c, out_c // 2, kernel_size=2, stride=2))
+            
+            self.ups.append(DoubleConv(out_c, out_c // 2, kernel_size=kernel_size, padding=padding))
+            
+            out_c //= 2
+
+        self.final_conv = nn.Conv2d(out_c, output_channels, kernel_size=1)
+
+    def forward(self, x):
+        skip_connections = []
+        
+        for down in self.downs:
+            x = down(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+
+        x = self.bottleneck(x)
+
+        for i in range(0, len(self.ups), 2):
+            x = self.ups[i](x)
+            skip_connection = skip_connections[-(i // 2 + 1)]
+            
+            if x.shape != skip_connection.shape:
+                x = nn.functional.interpolate(x, size=skip_connection.shape[2:], mode="bilinear", align_corners=True)
+                
+            x = torch.cat((skip_connection, x), dim=1)
+            x = self.ups[i + 1](x)
+
+        return self.final_conv(x)
+
+
+class UNet(BaseModel):
+    def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64, **kwargs):
+        super().__init__(name="UNet", **kwargs)
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        self.depth = depth
+        self.initial_filters = initial_filters
+
+    def create_model(self):
+        self.model = UnetTorch(
+            input_channels=self.input_channels,
+            output_channels=self.output_channels,
+            depth=self.depth,
+            initial_filters=self.initial_filters
+        )
+        return self.model

openmlcore-1.0.0/src/openmlcore/models/attUnets.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+from .base_model import BaseModel
+from .lib.convolutions import DoubleConv
+from .lib.attention import AttentionGate
+
+class AttUnetTorch(nn.Module):
+    def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64,
+                 kernel_size=3, padding=1):
+        super(AttUnetTorch, self).__init__()
+        self.downs          = nn.ModuleList()
+        self.ups            = nn.ModuleList()
+        self.attentionGates = nn.ModuleList()
+        self.pool           = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        in_c = input_channels
+        encoders_output_channels = []
+        out_c = initial_filters
+
+        for _ in range(depth):
+            self.downs.append(DoubleConv(in_c, out_c, kernel_size=kernel_size, padding=padding))
+            in_c = out_c
+            out_c *= 2
+            encoders_output_channels.append(in_c)
+
+        self.bottleneck = DoubleConv(in_c, out_c, kernel_size=kernel_size, padding=padding)
+        
+        for enc_c in reversed(encoders_output_channels):
+            dec_c = out_c //2
+            self.attentionGates.append(AttentionGate(enc_c, dec_c, enc_c//2))
+            self.ups.append(nn.ConvTranspose2d(out_c, dec_c, kernel_size=2, stride=2))
+            self.ups.append(DoubleConv(out_c, dec_c, kernel_size=kernel_size, padding=padding))
+
+            out_c = dec_c
+        self.final_conv = nn.Conv2d(out_c, output_channels, kernel_size=1)
+
+    def forward(self, x):
+        skip_connections = []
+        
+        for down in self.downs:
+            x = down(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+
+        x = self.bottleneck(x)
+
+        for i in range(0, len(self.ups), 2):
+            x = self.ups[i](x)
+            skip_connection = skip_connections[-(i // 2 + 1)]
+
+            if x.shape != skip_connection.shape:
+                x = nn.functional.interpolate(x, size=skip_connection.shape[2:], mode="bilinear", align_corners=True)
+
+            att_out = self.attentionGates[i // 2] (skip_connection, x)
+            x = torch.cat((att_out, x), dim=1)
+            x = self.ups[i + 1](x)
+
+        return self.final_conv(x)
+
+
+class AttentionUNet(BaseModel):
+    def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64, **kwargs):
+        super().__init__(name="AttentionUNet", **kwargs)
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        self.depth = depth
+        self.initial_filters = initial_filters
+
+    def create_model(self):
+        self.model = AttUnetTorch(
+            input_channels=self.input_channels,
+            output_channels=self.output_channels,
+            depth=self.depth,
+            initial_filters=self.initial_filters
+        )
+        return self.model

openmlcore-1.0.0/src/openmlcore/models/base_model.py

@@ -0,0 +1,227 @@
+import torch
+import torch.nn as nn
+from torchinfo import summary
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+import numpy as np
+
+class BaseModel(nn.Module):
+    def __init__(self, name: str = "BaseModel", model = None, metrics = None, loss_fn = None, optimizer = None, device = None):
+        super().__init__()
+        self.name = name
+        self.model = model
+        self.metrics = metrics
+        self.loss_fn = loss_fn
+        self.optimizer = optimizer
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    def help(self):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    def create_model(self):
+        raise NotImplementedError("Subclasses should implement this method.")
+        
+    def train(self, train_loader: DataLoader,val_loader:DataLoader | None = None, epochs: int = 100, save_checkpoint = True):
+        if self.name == "BaseModel":
+            raise NotImplementedError("Subclasses should implement this method.")
+        if self.model == None or self.name == "BaseModel":
+            raise ValueError("Model has not been created. Please call create_model() before training.")
+        if self.loss_fn == None:
+            raise ValueError("Loss function has not been set. Please set a loss function before training.")
+        if self.optimizer == None:
+            raise ValueError("Optimizer has not been set. Please set an optimizer before training.")
+
+        self.model.to(self.device)
+        for current_epoch in range(epochs):
+            # For each loop: a first phase of training and then validation.
+            # Training:
+            self.model.train()
+            total_loss = 0.0
+            total_steps = len(train_loader)
+            progress_bar = tqdm(
+                train_loader,
+                total=total_steps,
+                desc=f"Epoch {current_epoch + 1}/{epochs}",
+                unit="batch",
+                dynamic_ncols=True,
+            )
+            for idx, (X_batch, Y_batch) in enumerate(progress_bar, start=1):
+                X_batch, Y_batch = X_batch.to(self.device), Y_batch.to(self.device)
+                
+                self.optimizer.zero_grad()
+                prediction = self.model(X_batch)
+
+                loss = self.loss_fn(prediction, Y_batch)
+                loss.backward()
+                self.optimizer.step()
+                batch_loss = loss.item()
+                total_loss += batch_loss
+                progress_bar.set_postfix(
+                    batch_loss=f"{batch_loss:.4f}",
+                    avg_loss=f"{total_loss / idx:.4f}",
+                )
+            avg_loss = total_loss / len(train_loader)
+            tqdm.write(f"Epoch {current_epoch + 1}/{epochs} — average loss: {avg_loss:.4f}")
+
+            # Validation:
+            if val_loader:
+                self.model.eval()
+
+                dict_metrics_sum = {}
+                
+                val_progress_bar = tqdm(
+                    val_loader,
+                    total=len(val_loader),
+                    desc=f"Epoch {current_epoch + 1}/{epochs}",
+                    unit="batch",
+                    dynamic_ncols=True,
+                )
+
+                # Take a portion of the data and test:
+                with torch.no_grad():
+                    for idx, (X_batch, Y_batch) in enumerate(val_progress_bar, start=1):
+                        X_batch, Y_batch = X_batch.to(self.device), Y_batch.to(self.device)
+
+                        prediction = self.model(X_batch)
+
+                        for name_metric, metric_fn in self.metrics.items():
+                            metric_value = metric_fn(prediction, Y_batch)
+
+                            if isinstance(metric_value, torch.Tensor):
+                                metric_value = metric_value.detach().cpu()
+                                if metric_value.ndim > 1:
+                                    metric_value = metric_value.mean(dim=0)
+                                
+                                if metric_value.ndim == 0:
+                                    metric_value = metric_value.item()
+                                else:
+                                    metric_value = metric_value.tolist()
+
+                            elif isinstance(metric_value, np.ndarray):
+                                if metric_value.ndim > 1:
+                                    metric_value = metric_value.mean(axis=0)
+                                
+                                if metric_value.ndim == 0:
+                                    metric_value = metric_value.item()
+                                else:
+                                    metric_value = metric_value.tolist()
+                            if isinstance(metric_value, list) and len(metric_value) > 0 and isinstance(metric_value[0], list):
+                                metric_value = [sum(col) / len(col) for col in zip(*metric_value)]
+
+                            if name_metric not in dict_metrics_sum:
+                                if isinstance(metric_value, list):
+                                    dict_metrics_sum[name_metric] = [0.0] * len(metric_value)
+                                else:
+                                    dict_metrics_sum[name_metric] = 0.0
+
+                            if isinstance(metric_value, list):
+                                for cls_i, val in enumerate(metric_value):
+                                    dict_metrics_sum[name_metric][cls_i] += float(val)
+                            else:
+                                dict_metrics_sum[name_metric] += float(metric_value)
+
+                num_batchs = len(val_loader)
+                metrics_summary_str = []
+
+                for name_metric, val_sum in dict_metrics_sum.items():
+                    if isinstance(val_sum, list):
+                        avg_list = [v/num_batchs for v in val_sum]
+                        mean_all_classes = sum(avg_list)/len(avg_list)
+
+                        class_details = ", ".join([f"C{i}:{v:.3f}" for i, v in enumerate(avg_list)])
+                        metrics_summary_str.append(f"{name_metric}_Mean: {mean_all_classes:.4f} ({class_details})")
+                    else:
+                        avg_val = val_sum / num_batchs
+                        metrics_summary_str.append(f"{name_metric}: {avg_val:.4f}")
+
+                formatted_metrics = " | ".join(metrics_summary_str)
+                tqdm.write(f"Epoch {current_epoch + 1}/{epochs} — Validation: {formatted_metrics}")
+            
+            if save_checkpoint:
+                self.save_model(f"checkpoint_{self.name}.pt")
+
+    def predict(self, data):
+        if self.model == None or self.name == "BaseModel":
+            raise ValueError("Model has not been created. Please call create_model() before predicting.")
+
+        if data is not None:
+            data = data.to(self.device)
+            self.model.eval()
+            with torch.no_grad():
+                return self.model(data)
+        else:
+            raise ValueError("Data is None!")
+        
+
+    def evaluate(self, eval_loader: DataLoader):
+        if self.model == None or self.name == "BaseModel":
+            raise ValueError("Model has not been created. Please call create_model() before evaluating.")
+        if self.metrics is None:
+            raise ValueError("No metrics have been defined for evaluation. Please define metrics before evaluating.")
+
+        self.model.eval()
+        dict_metrics_sum = {}
+
+        with torch.no_grad():
+            for X_batch, Y_batch in eval_loader:
+                X_batch, Y_batch = X_batch.to(self.device), Y_batch.to(self.device)
+                predictions = self.model(X_batch)
+
+                for name_metric, metric_formula in self.metrics.items():
+                    metric_value = metric_formula(predictions, Y_batch) 
+                    
+                    if isinstance(metric_value, torch.Tensor):
+                        metric_value = metric_value.cpu().tolist()
+                    elif isinstance(metric_value, np.ndarray):
+                        metric_value = metric_value.tolist()
+
+                    if name_metric not in dict_metrics_sum:
+                        if isinstance(metric_value, list):
+                            dict_metrics_sum[name_metric] = [0.0] * len(metric_value)
+                        else:
+                            dict_metrics_sum[name_metric] = 0.0
+
+                    if isinstance(metric_value, list):
+                        for cls_i, val in enumerate(metric_value):
+                            dict_metrics_sum[name_metric][cls_i] += val
+                    else:
+                        dict_metrics_sum[name_metric] += metric_value
+            
+            num_batches = len(eval_loader)
+            final_metrics = {}
+
+            for name_metrics, val_sum in dict_metrics_sum.items():
+                if isinstance(val_sum, list):
+                    final_metrics[name_metric] = [v / num_batches for v in val_sum]
+                else:
+                    final_metrics[name_metric] = val_sum / num_batches
+
+        return final_metrics
+
+    def get_model_info(self):
+        if self.model == None or self.name == "BaseModel":
+            raise ValueError("Model has not been created. Please call create_model() before getting model info.")
+        
+        return summary(self.model)
+
+    def save_model(self, file_path):
+        self.model.eval()
+        torch.save(self.model.state_dict(), file_path)
+        print(f"Model saved to {file_path}")
+
+    def load_model(self, file_path, is_train_mode=False):
+        self.model.load_state_dict(torch.load(file_path, weights_only=True, map_location=self.device))
+        if is_train_mode:
+            self.model.train()
+        else:
+            self.model.eval()
+
+        # To make sure the modele is on the correct device:
+        self.model.to(self.device)
+        print(f"Model loaded from {file_path}, model is in {self.device} mode.")
+
+    def save_model_in_onnx(self, file_path, input_sample):
+        self.model.eval()
+        input_sample = input_sample.to(self.device)
+        torch.onnx.export(self.model, input_sample, file_path)
+        print(f"Model saved in ONNX format to {file_path}")

openmlcore-1.0.0/src/openmlcore/models/lib/attention.py

@@ -0,0 +1,21 @@
+import torch
+from torch import nn
+
+class AttentionGate(nn.Module):
+    def __init__(self, in_channels_encoder, in_channels_decoder, interchannel): 
+        super().__init__()
+        self.conv_enc = nn.Conv2d(in_channels_encoder, interchannel, kernel_size=1, stride=1, padding=0, bias=True)
+        self.conv_dec = nn.Conv2d(in_channels_decoder, interchannel, kernel_size=1, stride=1, padding=0, bias=True)
+        self.final_conv = nn.Conv2d(interchannel, 1, kernel_size=1, stride=1, padding=0, bias=True )
+
+        self.sigmoid = nn.Sigmoid()
+        self.relu    = nn.ReLU(inplace=True)  
+
+    def forward(self, encoder, decoder):
+        enc = self.conv_enc(encoder)
+        dec = self.conv_dec(decoder)
+        sum = self.relu(enc + dec)
+        x = self.sigmoid(self.final_conv(sum))
+        
+        return encoder * x
+

openmlcore-1.0.0/src/openmlcore/models/lib/convolutions.py

@@ -0,0 +1,15 @@
+import torch
+from torch import nn
+
+class DoubleConv(nn.Module):
+    def __init__(self, channel_in, channel_out, kernel_size = 3, padding = 1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels=channel_in , out_channels=channel_out, kernel_size=kernel_size, padding=padding)
+        self.conv2 = nn.Conv2d(in_channels=channel_out, out_channels=channel_out, kernel_size=kernel_size, padding=padding)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = nn.ReLU()(x)
+        x = self.conv2(x)
+        x = nn.ReLU()(x)
+        return x

openmlcore-1.0.0/src/openmlcore/models/unetpp.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .base_model import BaseModel
+from .lib.convolutions import DoubleConv
+
+class UnetPPTorch(nn.Module):
+    def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64,
+                 kernel_size=3, padding=1):
+        super(UnetPPTorch, self).__init__()
+        self.depth = depth
+        self.downs      = nn.ModuleList()
+        self.inter_down = nn.ModuleList()
+
+        filters = [initial_filters * (2**i) for i in range(self.depth + 1)]
+        self.pool = nn.MaxPool2d(kernel_size=2, stride= 2)
+        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+
+        self.convs = nn.ModuleDict()
+
+        for i in range(self.depth + 1):
+            for j in range (self.depth + 1 -i):
+                if j == 0:
+                    in_c = input_channels if i == 0 else filters[i-1]
+                    out_c = filters[i]
+
+                else:
+                    in_c = j * filters[i] + filters[i + 1]
+                    out_c = filters[i]
+
+                self.convs[f"conv_{i}_{j}"] = DoubleConv(in_c, out_c, kernel_size=kernel_size, padding=padding)
+
+        self.final_conv = DoubleConv(filters[0], output_channels, kernel_size=kernel_size, padding=padding)
+
+
+    def forward(self, x):
+        x_node = {}
+        x_node["0_0"] = self.convs["conv_0_0"](x)
+
+        # All the X_node on the left, the one called X^[i, 0]:
+        for i in range(1, self.depth + 1):
+            x_down = self.pool(x_node[f"{i-1}_0"])
+            x_node[f"{i}_0"] = self.convs[f"conv_{i}_0"](x_down)
+
+        # Loop for X^[i, j]
+        for j in range(1, self.depth +1):
+            for i in range(self.depth + 1 -j):
+                x_up = self.up(x_node[f"{i+1}_{j-1}"])
+
+                if x_up.shape[2:] != x_node[f'{i}_0'].shape[2:]:
+                    x_up = F.interpolate(
+                        x_up, size=x_node[f'{i}_0'].shape[2:], 
+                        mode="bilinear", align_corners=True
+                    )
+                
+                skip_concat = [x_node[f'{i}_{k}'] for k in range(j)]
+                skip_concat.append(x_up)
+                
+                concat_features = torch.cat(skip_concat, dim=1)
+                x_node[f'{i}_{j}'] = self.convs[f'conv_{i}_{j}'](concat_features)
+
+
+        # TODO: implementation of Deep supervision:
+        # if self.deepsupervision:
+        #     # TODO
+        # else:
+
+        return self.final_conv(x_node[f"0_{self.depth}"])
+
+    
+class UNetPP(BaseModel):
+    def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64, **kwargs):
+        super().__init__(name="UNet++", **kwargs)
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        self.depth = depth
+        self.initial_filters = initial_filters
+
+    def create_model(self):
+        self.model = UnetPPTorch(
+            input_channels=self.input_channels,
+            output_channels=self.output_channels,
+            depth=self.depth,
+            initial_filters=self.initial_filters
+        )
+        return self.model

openmlcore-1.0.0/src/openmlcore.egg-info/PKG-INFO

@@ -0,0 +1,76 @@
+Metadata-Version: 2.4
+Name: openmlcore
+Version: 1.0.0
+Summary: OpenMLCore is an open source framework (wrapper which simplify the use of PyTorch)
+Author-email: Doyez Sébastien <doyez.sebastien34090@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 Votre Nom
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+Project-URL: Homepage, https://sebastien-doyez2812.github.io/openMLCore.github.io/
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10.0
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: tqdm
+Requires-Dist: onnx
+Requires-Dist: onnxscript
+Requires-Dist: torchinfo
+Requires-Dist: datasets
+Dynamic: license-file
+
+# OpenML-core, Open source Package for Machine learning
+
+
+## Why OpenML?
+Today's computer vision ecosystem is heavily fragmented. If you want YOLO, Ultralytics makes it trivial. But the moment you need advanced semantic segmentation, hybrid transformer backbones, or efficient fine-tuning on modern architectures, you end up stitching together disjointed codebases, mismatched tensor shapes, and unmaintained repositories.
+
+OpenML bridges this gap by bringing a clean, Hugging Face-inspired Developer Experience (DX) to advanced computer vision tasks—from classical models to SOTA transformers and parameter-efficient fine-tuning.
+## Supported Architectures:
+
+- Classic & Deep Segmentation: U-Net, Attention U-Net, UNet++, Mask R-CNN
+
+- Advanced Transformers: Swin Transformer, TransUNet, MobileUNetR
+
+- SOTA & Edge: Vision Transformers, SAM 2 adaptations via efficient Low-Rank Adaptation (LoRA)
+## Key features:
+
+Unified Object-Oriented API: Every model inherits from a clean ModelBase class supporting standard .train(), .eval(), and .help() workflows.
+
+Native LoRA Wrappers: Fine-tune massive vision models and transformers on modest hardware (like a single RTX 3070) without blowing up your VRAM.
+
+Hugging Face Hub Integration: Push and pull checkpoints effortlessly using native .safetensors formatting.
+
+MLflow Tracking Out-of-the-Box: Keep tabs on your losses, metrics, and hyperparameter logs without writing boilerplate monitoring code.
+
+Interactive Help System: Running model.help() directly opens up the exact architectural documentation page on the web.
+
+## Contributions:
+
+Contributions, feature requests, and bug reports are welcome! Feel free to open an Issue or submit a Pull Request on GitHub.
+
+## Author
+
+Sébastien Doyez
+ML & AI Engineer,
+Interested in Agentics, Computer Vision, Robotics & Deep Learning

openmlcore-1.0.0/src/openmlcore.egg-info/SOURCES.txt

@@ -0,0 +1,21 @@
+LICENSE
+README.md
+pyproject.toml
+src/openmlcore/__init__.py
+src/openmlcore.egg-info/PKG-INFO
+src/openmlcore.egg-info/SOURCES.txt
+src/openmlcore.egg-info/dependency_links.txt
+src/openmlcore.egg-info/requires.txt
+src/openmlcore.egg-info/top_level.txt
+src/openmlcore/base/data.py
+src/openmlcore/loss/loss.py
+src/openmlcore/metrics/metrics.py
+src/openmlcore/models/UNets.py
+src/openmlcore/models/attUnets.py
+src/openmlcore/models/base_model.py
+src/openmlcore/models/unetpp.py
+src/openmlcore/models/lib/attention.py
+src/openmlcore/models/lib/convolutions.py
+tests/test_base_model.py
+tests/test_loss.py
+tests/test_main.py

openmlcore-1.0.0/src/openmlcore.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

openmlcore-1.0.0/src/openmlcore.egg-info/requires.txt

@@ -0,0 +1,6 @@
+torch
+tqdm
+onnx
+onnxscript
+torchinfo
+datasets

openmlcore-1.0.0/src/openmlcore.egg-info/top_level.txt

@@ -0,0 +1 @@
+openmlcore

openmlcore-1.0.0/tests/test_base_model.py

@@ -0,0 +1,101 @@
+import pytest
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchgen import model
+from torchgen import model
+from openmlcore.models.base_model import *
+from openmlcore.metrics.metrics import *
+
+# class DummyModel(BaseModel):
+#     def __init__(self):
+#         super().__init__(name = "DummyLinearModel")
+#         self.create_model()
+#         self.loss_fn = nn.CrossEntropyLoss()
+#         self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01)
+#         self.metrics = [("MSE", lambda pred, target: nn.functional.mse_loss(pred, target))]
+
+#     def help(self):
+#         print("This is a dummy linear model for testing purposes.")
+
+#     def create_model(self):
+#         self.fc = nn.Linear(10, 2)
+#         self.model = self.fc
+
+class DummyModel(BaseModel):
+    def __init__(self):
+        super().__init__(name="DummyLinearModel")
+        # On appelle super().__init__() avant d'assigner self.model
+        self.create_model()
+        self.loss_fn = nn.CrossEntropyLoss()
+        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01)
+        self.metrics = {
+            "Accuracy": accuracy,
+            "IoU": iou
+        }
+        
+    def help(self):
+        print("This is a dummy linear model for testing purposes.")
+
+    def create_model(self):
+        # Utiliser nn.Sequential garantit que tous les sous-modules 
+        # basculent sur le bon device en même temps que self.model
+        self.model = nn.Sequential(
+            nn.Linear(10, 2)
+        ).to(self.device)
+
+@pytest.fixture
+def dummy_dataloader():
+    X = torch.randn(64, 10)
+    Y = torch.randn(64, 2)
+    dataset = torch.utils.data.TensorDataset(X, Y)
+    return DataLoader(dataset, batch_size=8, shuffle=True)
+
+
+def test_abstract_base_modeel():
+    base_model = BaseModel()
+    with pytest.raises(NotImplementedError):
+        base_model.create_model()
+
+    with pytest.raises(NotImplementedError):    
+        base_model.help()
+
+def test_training_exceptions():
+    base_model = BaseModel()
+    loader = torch.utils.data.DataLoader(torch.randn(64, 10), batch_size=8)
+
+    with pytest.raises(NotImplementedError):
+        base_model.train(loader, epochs=1)
+
+def test_model_training(dummy_dataloader):
+    model = DummyModel()
+    model.train(dummy_dataloader, epochs=2)
+
+def test_model_prediction(dummy_dataloader):
+    model = DummyModel()
+    print(model.device)
+    X_sample = torch.randn(5, 10).to(model.device)
+    preds = model.predict(X_sample)
+    assert preds.shape == (5, 2)
+
+def test_model_save_and_load(tmp_path, dummy_dataloader):
+    model = DummyModel()
+    file_path = tmp_path / "model.pt"
+    
+    model.save_model(str(file_path))
+    assert file_path.exists()
+
+    new_model = DummyModel()
+    new_model.load_model(str(file_path), is_train_mode=False)
+    
+    for p1, p2 in zip(model.model.parameters(), new_model.model.parameters()):
+        assert torch.equal(p1, p2)
+
+
+def test_onnx_export(tmp_path):
+    model = DummyModel()
+    file_path = tmp_path / "model.onnx"
+    dummy_input = torch.randn(1, 10).to(model.device)
+    
+    model.save_model_in_onnx(str(file_path), dummy_input)
+    assert file_path.exists()    

openmlcore-1.0.0/tests/test_loss.py

@@ -0,0 +1,27 @@
+from openmlcore.loss.loss import *
+
+def test_BCEWithLogits_loss():
+    loss = BCEWithLogitsLoss()
+    assert loss is not None
+
+def test_mean_squared_error_loss():
+    loss = MeanSquaredErrorLoss()
+    assert loss is not None
+
+def test_dice_loss():
+    loss = DiceLoss()
+    assert loss is not None
+
+def test_focal_loss():
+    loss = FocalLoss()
+    assert loss is not None
+
+def test_huber_loss():
+    loss = HuberLoss()
+    assert loss is not None
+
+def test_kl_divergence_loss():
+    loss = KLDivergenceLoss()
+    assert loss is not None
+
+    

openmlcore-1.0.0/tests/test_main.py

@@ -0,0 +1,4 @@
+# TODO: implement tests here...
+
+def test_example():
+    assert 1 + 1 == 2