openmlcore 1.0.0__tar.gz → 1.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (26) hide show
  1. {openmlcore-1.0.0/src/openmlcore.egg-info → openmlcore-1.2.0}/PKG-INFO +1 -24
  2. {openmlcore-1.0.0 → openmlcore-1.2.0}/pyproject.toml +1 -1
  3. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/loss/loss.py +53 -71
  4. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/metrics/metrics.py +3 -2
  5. openmlcore-1.2.0/src/openmlcore/models/MobileUnetR.py +329 -0
  6. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/models/base_model.py +54 -5
  7. openmlcore-1.2.0/src/openmlcore/models/segformer.py +322 -0
  8. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/models/unetpp.py +1 -1
  9. {openmlcore-1.0.0 → openmlcore-1.2.0/src/openmlcore.egg-info}/PKG-INFO +1 -24
  10. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore.egg-info/SOURCES.txt +2 -1
  11. {openmlcore-1.0.0 → openmlcore-1.2.0}/tests/test_base_model.py +8 -2
  12. openmlcore-1.2.0/tests/test_loss.py +26 -0
  13. openmlcore-1.0.0/LICENSE +0 -21
  14. openmlcore-1.0.0/tests/test_loss.py +0 -27
  15. {openmlcore-1.0.0 → openmlcore-1.2.0}/README.md +0 -0
  16. {openmlcore-1.0.0 → openmlcore-1.2.0}/setup.cfg +0 -0
  17. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/__init__.py +0 -0
  18. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/base/data.py +0 -0
  19. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/models/UNets.py +0 -0
  20. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/models/attUnets.py +0 -0
  21. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/models/lib/attention.py +0 -0
  22. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore/models/lib/convolutions.py +0 -0
  23. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore.egg-info/dependency_links.txt +0 -0
  24. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore.egg-info/requires.txt +0 -0
  25. {openmlcore-1.0.0 → openmlcore-1.2.0}/src/openmlcore.egg-info/top_level.txt +0 -0
  26. {openmlcore-1.0.0 → openmlcore-1.2.0}/tests/test_main.py +0 -0
@@ -1,43 +1,20 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: openmlcore
3
- Version: 1.0.0
3
+ Version: 1.2.0
4
4
  Summary: OpenMLCore is an open source framework (wrapper which simplify the use of PyTorch)
5
5
  Author-email: Doyez Sébastien <doyez.sebastien34090@gmail.com>
6
- License: MIT License
7
-
8
- Copyright (c) 2026 Votre Nom
9
-
10
- Permission is hereby granted, free of charge, to any person obtaining a copy
11
- of this software and associated documentation files (the "Software"), to deal
12
- in the Software without restriction, including without limitation the rights
13
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
- copies of the Software, and to permit persons to whom the Software is
15
- furnished to do so, subject to the following conditions:
16
-
17
- The above copyright notice and this permission notice shall be included in all
18
- copies or substantial portions of the Software.
19
-
20
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
- SOFTWARE.
27
6
  Project-URL: Homepage, https://sebastien-doyez2812.github.io/openMLCore.github.io/
28
7
  Classifier: Programming Language :: Python :: 3
29
8
  Classifier: License :: OSI Approved :: MIT License
30
9
  Classifier: Operating System :: OS Independent
31
10
  Requires-Python: >=3.10.0
32
11
  Description-Content-Type: text/markdown
33
- License-File: LICENSE
34
12
  Requires-Dist: torch
35
13
  Requires-Dist: tqdm
36
14
  Requires-Dist: onnx
37
15
  Requires-Dist: onnxscript
38
16
  Requires-Dist: torchinfo
39
17
  Requires-Dist: datasets
40
- Dynamic: license-file
41
18
 
42
19
  # OpenML-core, Open source Package for Machine learning
43
20
 
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "openmlcore"
7
- version = "1.0.0"
7
+ version = "1.2.0"
8
8
  authors = [
9
9
  { name="Doyez Sébastien", email="doyez.sebastien34090@gmail.com" },
10
10
  ]
@@ -29,6 +29,9 @@ class BaseLoss(nn.Module):
29
29
  class BCEWithLogitsLoss(BaseLoss):
30
30
  def __init__(self, weights=None):
31
31
  super().__init__(name="BCEWithLogitsLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html")
32
+ if weights is None:
33
+ raise ValueError("weights need to be set!")
34
+
32
35
  device = "cuda" if torch.cuda.is_available() else "cpu"
33
36
 
34
37
  pos_weights = None
@@ -48,6 +51,9 @@ class BCEWithLogitsLoss(BaseLoss):
48
51
  class DiceLoss(BaseLoss):
49
52
  def __init__(self, weights = None):
50
53
  super().__init__(name="DiceLoss", doc_url="https://arxiv.org/abs/1606.04797")
54
+ if weights is None:
55
+ raise ValueError("weights need to be set!")
56
+
51
57
  self.weights = weights
52
58
 
53
59
  def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
@@ -71,76 +77,67 @@ class DiceLoss(BaseLoss):
71
77
 
72
78
  return dice_loss.sum(dim=1).mean()
73
79
  return dice_loss.mean()
74
-
75
- # TODO:
76
- # 1 Hot encoding loss
77
- #
78
-
79
80
 
80
81
  class MeanSquaredErrorLoss(BaseLoss):
81
- def __init__(self, index_ignore=255):
82
+ def __init__(self, weights = None):
82
83
  super().__init__(name="MeanSquaredErrorLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html")
83
- self.index_ignore = index_ignore
84
+ if weights is None:
85
+ raise ValueError("weights need to be set!")
86
+
87
+ self.weights = weights
84
88
 
85
89
  def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
86
- num_classes = outputs.shape[1]
87
- probs = F.softmax(outputs, dim=1)
88
- targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
89
-
90
- mse = (probs - targets_one_hot) ** 2
91
- mse = mse * valid_mask
92
-
93
- num_valids = valid_mask.sum() * num_classes
94
- return mse.sum() / torch.clamp(num_valids, min=1.0)
90
+ probs = torch.sigmoid(outputs)
91
+ probs = probs.view (probs.size(0) , probs.size(1), -1)
92
+ targets = targets.view(probs.size(0), probs.size(1), -1)
95
93
 
94
+ mse = (probs - targets) ** 2
95
+ weighted_mse = 0
96
+ for idx, weight in enumerate(self.weights):
97
+ weighted_mse+= weight * mse[:, idx, :]
96
98
 
99
+ return weighted_mse.mean()
97
100
 
98
101
  class FocalLoss(BaseLoss):
99
- def __init__(self, alpha=0.25, gamma=2.0, index_ignore=255):
102
+ def __init__(self, alpha=0.25, gamma=2.0, weights = None):
100
103
  super().__init__(name="FocalLoss", doc_url="https://arxiv.org/abs/1708.02002")
104
+ if weights is None:
105
+ raise ValueError("weights need to be set!")
106
+
101
107
  self.alpha = alpha
102
108
  self.gamma = gamma
103
- self.index_ignore = index_ignore
109
+ self.weights = weights
104
110
 
105
111
  def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
106
- num_classes = outputs.shape[1]
107
- targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
108
-
109
- bce_loss = F.binary_cross_entropy_with_logits(outputs, targets_one_hot, reduction='none')
112
+ probs = outputs.view(outputs.size(0) , outputs.size(1), -1)
113
+ targets = targets.view(outputs.size(0), outputs.size(1), -1)
114
+
115
+ bce_loss = F.binary_cross_entropy_with_logits(probs, targets, reduction='none') # Apply a sigmoid inside...
110
116
  pt = torch.exp(-bce_loss)
111
117
  focal_loss = self.alpha * ((1 - pt) ** self.gamma) * bce_loss
112
- focal_loss = focal_loss * valid_mask
113
- num_valids = valid_mask.sum() * num_classes
114
- return focal_loss.sum() / torch.clamp(num_valids, min=1.0)
115
-
116
- class MeanSquaredErrorLoss(BaseLoss):
117
- def __init__(self, index_ignore=255):
118
- super().__init__(name="MeanSquaredErrorLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html")
119
- self.index_ignore = index_ignore
120
-
121
- def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
122
- num_classes = outputs.shape[1]
123
- probs = F.softmax(outputs, dim=1)
124
- targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
125
-
126
- mse = (probs - targets_one_hot) ** 2
127
- mse = mse * valid_mask
128
-
129
- num_valids = valid_mask.sum() * num_classes
130
- return mse.sum() / torch.clamp(num_valids, min=1.0)
131
118
 
119
+ final_loss = 0
120
+ for idx, weight in enumerate(self.weights):
121
+ final_loss += weight * focal_loss[:, idx, :]
122
+ return final_loss.mean()
123
+
132
124
  class HuberLoss(BaseLoss):
133
- def __init__(self, delta=1.0, index_ignore=255):
125
+ def __init__(self, delta=1.0, weights = None):
134
126
  super().__init__(name="HuberLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.HuberLoss.html")
127
+ if weights is None:
128
+ raise ValueError("weights need to be set!")
129
+ self.weights = weights
135
130
  self.delta = delta
136
- self.index_ignore = index_ignore
137
131
 
138
132
  def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
139
- num_classes = outputs.shape[1]
140
- probs = F.softmax(outputs, dim=1)
141
- targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
142
-
143
- error = probs - targets_one_hot
133
+ # num_classes = outputs.shape[1]
134
+ # probs = F.softmax(outputs, dim=1)
135
+
136
+ probs = torch.sigmoid(outputs)
137
+ probs = probs.view (probs.size(0) , probs.size(1), -1)
138
+ targets = targets.view(probs.size(0), probs.size(1), -1)
139
+
140
+ error = probs - targets
144
141
  abs_error = torch.abs(error)
145
142
 
146
143
  huber = torch.where(
@@ -148,30 +145,15 @@ class HuberLoss(BaseLoss):
148
145
  0.5 * (error ** 2),
149
146
  self.delta * (abs_error - 0.5 * self.delta)
150
147
  )
151
- huber = huber * valid_mask
152
-
153
- num_valids = valid_mask.sum() * num_classes
154
- return huber.sum() / torch.clamp(num_valids, min=1.0)
155
-
156
- class KLDivergenceLoss(BaseLoss):
157
- def __init__(self, index_ignore=255):
158
- super().__init__(name="KLDivergenceLoss", doc_url="https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html")
159
- self.eps = 1e-8
160
- self.index_ignore = index_ignore
161
-
162
- def forward(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
163
- num_classes = outputs.shape[1]
164
- log_probs = F.log_softmax(outputs, dim=1)
165
- targets_one_hot, valid_mask = self._to_one_hot_clean(targets, num_classes, self.index_ignore)
166
-
167
- # KL-Div = p(x) * (log(p(x)) - log(q(x)))
168
- # Pour une distribution One-Hot, cela se réduit à : - log(q(x)) sur la classe vraie
169
- kl = F.kl_div(log_probs, targets_one_hot, reduction='none')
170
- kl = kl * valid_mask
171
-
172
- num_valids = valid_mask.sum() * num_classes
173
- return kl.sum() / torch.clamp(num_valids, min=1.0)
174
148
 
149
+ if not isinstance(self.weights, torch.Tensor):
150
+ weights_tensor = torch.tensor(self.weights, device=outputs.device, dtype=outputs.dtype)
151
+ else:
152
+ weights_tensor = self.weights.to(outputs.device)
153
+ huber = huber * weights_tensor.view(1, -1, 1)
154
+
155
+ return huber.mean()
156
+
175
157
  ####################################
176
158
  # Custom Loss implementation #
177
159
  ####################################
@@ -34,8 +34,9 @@ def accuracy(y_pred, y_true):
34
34
  final_metric[idx] = metric / num_batches
35
35
  return final_metric
36
36
 
37
- def iou(y_pred, y_true,smooth = 1e-6):
38
- preds = (y_pred > 0.0).float()
37
+ def iou(y_pred, y_true, threshold = 0.5, smooth = 1e-6):
38
+ preds = (y_pred > threshold).float()
39
+
39
40
  preds = preds.view(preds.size(0), preds.size(1), -1)
40
41
  y_true = y_true.view(y_true.size(0), y_true.size(1), -1)
41
42
 
@@ -0,0 +1,329 @@
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+ import timm
5
+ from .base_model import BaseModel
6
+
7
+ class MobileNetBlockv2(nn.Module):
8
+ def __init__(self, in_channels, out_channels, stride=1, expand_ratio=4, **kwargs):
9
+ super().__init__(**kwargs)
10
+ self.stride = stride
11
+ self.use_residual = (self.stride == 1 and in_channels == out_channels)
12
+
13
+ hidden_dim = in_channels * expand_ratio
14
+ self.layers = nn.ModuleList()
15
+
16
+ if expand_ratio != 1:
17
+ self.layers.append(nn.Conv2d(in_channels, hidden_dim, kernel_size=1, bias=False))
18
+ self.layers.append(nn.BatchNorm2d(hidden_dim))
19
+ self.layers.append(nn.ReLU6(inplace=True)) # ReLU6 est souvent préféré pour MobileNet
20
+ else:
21
+ hidden_dim = in_channels
22
+
23
+ self.layers.append(nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=self.stride, padding=1, groups=hidden_dim, bias=False))
24
+ self.layers.append(nn.BatchNorm2d(hidden_dim))
25
+ self.layers.append(nn.ReLU6(inplace=True))
26
+
27
+ self.layers.append(nn.Conv2d(hidden_dim, out_channels, kernel_size=1, bias=False))
28
+ self.layers.append(nn.BatchNorm2d(out_channels))
29
+
30
+ def forward(self, X):
31
+ inputX = X if self.use_residual else None
32
+
33
+ for layer in self.layers:
34
+ X = layer(X)
35
+
36
+ if inputX is not None:
37
+ return inputX + X
38
+ return X
39
+
40
+ class STEM(nn.Module):
41
+ def __init__(self, in_channels=3, **kwargs):
42
+ super().__init__(**kwargs)
43
+
44
+ self.conv2D = nn.Conv2d(in_channels=in_channels, out_channels=32, kernel_size=3, stride=2, padding=1, bias=False)
45
+ self.batchNorm = nn.BatchNorm2d(32)
46
+ self.relu = nn.ReLU6(inplace=True)
47
+
48
+ self.mobile_block1 = MobileNetBlockv2(in_channels=32, out_channels=16, stride=1, expand_ratio=1)
49
+ self.mobile_block2a = MobileNetBlockv2(in_channels=16, out_channels=24, stride=2, expand_ratio=6)
50
+ self.mobile_block2b = MobileNetBlockv2(in_channels=24, out_channels=24, stride=1, expand_ratio=6)
51
+ self.mobile_block3 = MobileNetBlockv2(in_channels=24, out_channels=24, stride=1, expand_ratio=6)
52
+
53
+ def forward(self, X):
54
+ X = self.relu(self.batchNorm(self.conv2D(X)))
55
+ X = self.mobile_block1(X)
56
+ X = self.mobile_block2a(X)
57
+ X = self.mobile_block2b(X)
58
+ X = self.mobile_block3(X)
59
+ return X
60
+
61
+ class MobileViTBlock(nn.Module):
62
+ def __init__(self, in_channels, out_channels, patch_size=2, l=2, dim=64, **kwargs):
63
+ super().__init__(**kwargs)
64
+ self.patchsize = patch_size
65
+ self.dim = dim
66
+ self.out_channels = out_channels
67
+ self.L = l
68
+
69
+ self.conv3x3_1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, padding=1, bias=False)
70
+ self.conv1x1_1 = nn.Conv2d(in_channels=in_channels, out_channels=dim, kernel_size=1, padding=0, bias=False)
71
+
72
+ self.conv3x3_2 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1, bias=False)
73
+ self.conv1x1_2 = nn.Conv2d(in_channels=dim, out_channels=in_channels, kernel_size=1, padding=0, bias=False)
74
+
75
+ self.transformers = nn.ModuleList([ViTransformer(dim=self.patchsize * self.patchsize * dim) for _ in range(self.L)])
76
+
77
+ self.unfold = nn.Unfold(
78
+ kernel_size=patch_size,
79
+ stride=patch_size,
80
+ padding=0
81
+ )
82
+
83
+ def forward(self, X):
84
+ skip_connection = X
85
+
86
+ X = self.conv3x3_1(X)
87
+ X = self.conv1x1_1(X)
88
+
89
+ B, C, H, W = X.shape
90
+ pH, pW = H // self.patchsize, W // self.patchsize
91
+
92
+ X = self.unfold(X)
93
+ X = X.reshape(B, C * self.patchsize * self.patchsize, pH * pW)
94
+ X = X.permute(0, 2, 1)
95
+
96
+ for transformer in self.transformers:
97
+ X = transformer(X)
98
+
99
+ X = X.permute(0, 2, 1)
100
+ X = X.reshape(B, C, self.patchsize, self.patchsize, pH, pW)
101
+ X = X.permute(0, 1, 4, 2, 5, 3)
102
+ X = X.reshape(B, C, pH * self.patchsize, pW * self.patchsize)
103
+
104
+ X = self.conv1x1_2(X)
105
+ X = X + skip_connection
106
+ X = self.conv3x3_2(X)
107
+
108
+ return X
109
+
110
+ class ViTransformer(nn.Module):
111
+ def __init__(self, dim, num_heads=8, ffn_ratio=4, dropout=0.0, **kwargs):
112
+ super().__init__(**kwargs)
113
+ self.dim = dim
114
+ self.num_heads = num_heads
115
+ self.ffn_ratio = ffn_ratio
116
+ self.dropout_rate = dropout
117
+
118
+ self.norm1 = nn.LayerNorm(dim, eps=1e-6)
119
+ self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=num_heads, batch_first=True, dropout=dropout)
120
+ self.drop1 = nn.Dropout(dropout)
121
+
122
+ self.norm2 = nn.LayerNorm(dim, eps=1e-6)
123
+ self.ffn_1 = nn.Linear(in_features=dim, out_features=dim * ffn_ratio)
124
+ self.gelu = nn.GELU()
125
+ self.ffn_2 = nn.Linear(in_features=dim * ffn_ratio, out_features=dim)
126
+ self.drop2 = nn.Dropout(dropout)
127
+
128
+ def forward(self, X):
129
+ input_tensor = X
130
+ X = self.norm1(X)
131
+ X, _ = self.attn(X, X, X)
132
+ X = self.drop1(X)
133
+ X = X + input_tensor
134
+
135
+ res = X
136
+ X = self.norm2(X)
137
+ X = self.ffn_1(X)
138
+ X = self.gelu(X)
139
+ X = self.drop2(X)
140
+ X = self.ffn_2(X)
141
+ X = self.drop2(X)
142
+ return X + res
143
+
144
+ class HybridEncoderBlock(nn.Module):
145
+ def __init__(self, input_channels, out_channels, stride=2, **kwargs):
146
+ super().__init__(**kwargs)
147
+ self.Conv1x1 = nn.Conv2d(in_channels=input_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=False)
148
+ self.Conv3x3 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
149
+ self.mobileViT_block = MobileViTBlock(in_channels=out_channels, out_channels=out_channels)
150
+
151
+ def forward(self, X):
152
+ X = self.Conv1x1(X)
153
+ X = self.Conv3x3(X)
154
+ X = self.mobileViT_block(X)
155
+ return X
156
+
157
+ class HybridDecoderBlock(nn.Module):
158
+ def __init__(self, in_channels, skip_channels, out_channels, **kwargs):
159
+ super().__init__(**kwargs)
160
+ self.transpose_conv = nn.ConvTranspose2d(
161
+ in_channels=in_channels,
162
+ out_channels=out_channels,
163
+ kernel_size=2,
164
+ stride=2
165
+ )
166
+
167
+ concat_channels = out_channels + skip_channels
168
+
169
+ self.conv1x1 = nn.Conv2d(
170
+ in_channels=concat_channels,
171
+ out_channels=out_channels,
172
+ kernel_size=1,
173
+ stride=1,
174
+ bias=False
175
+ )
176
+ self.silu1 = nn.SiLU()
177
+
178
+ self.conv_depth_wise = nn.Conv2d(
179
+ in_channels=out_channels,
180
+ out_channels=out_channels,
181
+ kernel_size=3,
182
+ stride=1,
183
+ padding=1,
184
+ groups=out_channels,
185
+ bias=False
186
+ )
187
+ self.silu2 = nn.SiLU()
188
+
189
+ self.conv1x1_2 = nn.Conv2d(
190
+ in_channels=out_channels,
191
+ out_channels=out_channels,
192
+ kernel_size=1,
193
+ stride=1,
194
+ bias=False
195
+ )
196
+ self.silu3 = nn.SiLU()
197
+
198
+ self.mobileViTBlock = MobileViTBlock(
199
+ in_channels=out_channels,
200
+ out_channels=out_channels
201
+ )
202
+
203
+ def forward(self, decoder_in, skip_in):
204
+ x = self.transpose_conv(decoder_in)
205
+
206
+ if x.shape[2:] != skip_in.shape[2:]:
207
+ x = F.interpolate(x, size=skip_in.shape[2:], mode='bilinear', align_corners=False)
208
+
209
+ x = torch.cat([x, skip_in], dim=1)
210
+ x = self.silu1(self.conv1x1(x))
211
+ x = self.silu2(self.conv_depth_wise(x))
212
+ x = self.silu3(self.conv1x1_2(x))
213
+ return self.mobileViTBlock(x)
214
+
215
+ class MobileUnetRTorch(nn.Module):
216
+ def __init__(self, input_channels, output_channels):
217
+ super().__init__()
218
+ self.STEM = STEM(in_channels=input_channels) #should be 3, because RGB = 3 channels
219
+ # HEB = Hybrid Encoder Block
220
+ self.HEB_1 = HybridEncoderBlock(input_channels=24, out_channels=64, stride=2) # H/8
221
+ self.HEB_2 = HybridEncoderBlock(input_channels=64, out_channels=128, stride=2) # H/16
222
+ self.HEB_3 = HybridEncoderBlock(input_channels=128, out_channels=256, stride=2) # H/32
223
+ self.HEB_4 = HybridEncoderBlock(input_channels=256, out_channels=512, stride=2) # H/64
224
+
225
+ # HDB = Hybrid Decoder Block
226
+ self.HDB_1 = HybridDecoderBlock(in_channels= 512, skip_channels=256, out_channels=256) #H/32
227
+ self.HDB_2 = HybridDecoderBlock(in_channels= 256, skip_channels=128, out_channels=128) #H/16
228
+ self.HDB_3 = HybridDecoderBlock(in_channels= 128, skip_channels=64, out_channels=64) #H/8
229
+
230
+ self.final_conv = nn.ConvTranspose2d(in_channels=64, out_channels=output_channels, kernel_size=4, stride=4, padding=0)
231
+ self.sigmoid = nn.Sigmoid()
232
+
233
+
234
+ def forward(self, X):
235
+ input_size = X.shape[2:]
236
+ skip_connection = []
237
+
238
+ X = self.STEM(X) # Embedding
239
+ skip_connection.append(X)
240
+ X = self.HEB_1(X)
241
+ skip_connection.append(X)
242
+
243
+ X = self.HEB_2(X)
244
+ skip_connection.append(X)
245
+
246
+ X = self.HEB_3(X)
247
+ skip_connection.append(X)
248
+
249
+ X = self.HEB_4(X)
250
+
251
+ X = self.HDB_1(X, skip_connection.pop())
252
+ X = self.HDB_2(X, skip_connection.pop())
253
+ X = self.HDB_3(X, skip_connection.pop())
254
+
255
+ X = self.final_conv(X)
256
+
257
+ if X.shape[2:] != input_size:
258
+ X = F.interpolate(X, size=input_size, mode='bilinear', align_corners=False)
259
+ return X
260
+ # return self.sigmoid(X)
261
+
262
+ class MobileUnetRTorch_v2(nn.Module):
263
+ def __init__(self, input_channels, output_channels):
264
+ super().__init__()
265
+
266
+ self.encoder = timm.create_model(
267
+ "mobilevit_s.cvnets_in1k",
268
+ pretrained=True,
269
+ features_only=True,
270
+ out_indices=(1, 2, 3, 4)
271
+ )
272
+ for param in self.encoder.parameters():
273
+ param.requires_grad = True
274
+
275
+ # features[3] (bottleneck) = 640 channels
276
+ # features[2] (skip3) = 128 channels
277
+ # features[1] (skip2) = 96 channels
278
+ # features[0] (skip1) = 64 channels
279
+
280
+ self.HDB_1 = HybridDecoderBlock(in_channels=640, skip_channels=128, out_channels=256)
281
+ self.HDB_2 = HybridDecoderBlock(in_channels=256, skip_channels=96, out_channels=128)
282
+ self.HDB_3 = HybridDecoderBlock(in_channels=128, skip_channels=64, out_channels=64)
283
+
284
+ self.final_conv = nn.ConvTranspose2d(in_channels=64, out_channels=output_channels, kernel_size=4, stride=4, padding=0)
285
+ self.sigmoid = nn.Sigmoid()
286
+
287
+ def forward(self, X):
288
+ input_size = X.shape[2:]
289
+
290
+ features = self.encoder(X)
291
+ skip1, skip2, skip3, bottleneck = features[0], features[1], features[2], features[3]
292
+
293
+ X = self.HDB_1(bottleneck, skip3)
294
+ X = self.HDB_2(X, skip2)
295
+ X = self.HDB_3(X, skip1)
296
+
297
+ X = self.final_conv(X)
298
+
299
+ if X.shape[2:] != input_size:
300
+ X = F.interpolate(X, size=input_size, mode='bilinear', align_corners=False)
301
+
302
+ return X
303
+ # return self.sigmoid(X)
304
+
305
+ class MobileUnetR(BaseModel):
306
+ def __init__(self, input_channels=1, output_channels=1, **kwargs):
307
+ super().__init__(name="MobileUnetR", **kwargs)
308
+ self.input_channels = input_channels
309
+ self.output_channels = output_channels
310
+
311
+ def create_model(self):
312
+ self.model = MobileUnetRTorch(
313
+ input_channels=self.input_channels,
314
+ output_channels=self.output_channels
315
+ )
316
+ return self.model
317
+
318
+ class MobileUnetR_v2(BaseModel):
319
+ def __init__(self, input_channels=1, output_channels=1, **kwargs):
320
+ super().__init__(name="MobileUnetR_apple_version", **kwargs)
321
+ self.input_channels = input_channels
322
+ self.output_channels = output_channels
323
+
324
+ def create_model(self):
325
+ self.model = MobileUnetRTorch_v2(
326
+ input_channels=self.input_channels,
327
+ output_channels=self.output_channels
328
+ )
329
+ return self.model
@@ -1,4 +1,6 @@
1
1
  import torch
2
+ import csv
3
+ import mlflow
2
4
  import torch.nn as nn
3
5
  from torchinfo import summary
4
6
  from torch.utils.data import DataLoader
@@ -6,7 +8,7 @@ from tqdm.auto import tqdm
6
8
  import numpy as np
7
9
 
8
10
  class BaseModel(nn.Module):
9
- def __init__(self, name: str = "BaseModel", model = None, metrics = None, loss_fn = None, optimizer = None, device = None):
11
+ def __init__(self, name: str = "BaseModel", model = None, metrics = None, loss_fn = None, optimizer = None, device = None, write_result_in_CSV = True, csv_path= None, use_mlflow = True):
10
12
  super().__init__()
11
13
  self.name = name
12
14
  self.model = model
@@ -14,6 +16,13 @@ class BaseModel(nn.Module):
14
16
  self.loss_fn = loss_fn
15
17
  self.optimizer = optimizer
16
18
  self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
19
+
20
+ self.use_mlflow = use_mlflow
21
+ self.write_result_in_CSV = write_result_in_CSV
22
+ if csv_path is None:
23
+ self.csv_path = f"{self.name}.csv"
24
+ else:
25
+ self.csv_path = csv_path
17
26
 
18
27
  def help(self):
19
28
  raise NotImplementedError("Subclasses should implement this method.")
@@ -32,6 +41,27 @@ class BaseModel(nn.Module):
32
41
  raise ValueError("Optimizer has not been set. Please set an optimizer before training.")
33
42
 
34
43
  self.model.to(self.device)
44
+
45
+ # MLFlow:
46
+ if self.use_mlflow:
47
+ lr = self.optimizer.param_groups[0]['lr'] if self.optimizer and len(self.optimizer.param_groups) > 0 else None
48
+ mlflow.log_params({
49
+ "model_name": self.name,
50
+ "epochs": epochs,
51
+ "learning_rate": lr,
52
+ "device": str(self.device)
53
+ })
54
+
55
+ # CSV Logger:
56
+ if self.write_result_in_CSV:
57
+ csv_fields = ["epoch", "train_loss"]
58
+ for name_metric in self.metrics.keys():
59
+ csv_fields.append(f"val_{name_metric}")
60
+
61
+ with open(self.csv_path, 'w', newline='', encoding='utf-8') as file:
62
+ writer = csv.DictWriter(file, fieldnames=csv_fields, extrasaction='ignore')
63
+ writer.writeheader()
64
+
35
65
  for current_epoch in range(epochs):
36
66
  # For each loop: a first phase of training and then validation.
37
67
  # Training:
@@ -63,6 +93,11 @@ class BaseModel(nn.Module):
63
93
  avg_loss = total_loss / len(train_loader)
64
94
  tqdm.write(f"Epoch {current_epoch + 1}/{epochs} — average loss: {avg_loss:.4f}")
65
95
 
96
+
97
+ # Create a dictionnary for the csv:
98
+ epoch_metrics = {"epoch": current_epoch,
99
+ "train_loss" : avg_loss}
100
+
66
101
  # Validation:
67
102
  if val_loader:
68
103
  self.model.eval()
@@ -130,13 +165,27 @@ class BaseModel(nn.Module):
130
165
 
131
166
  class_details = ", ".join([f"C{i}:{v:.3f}" for i, v in enumerate(avg_list)])
132
167
  metrics_summary_str.append(f"{name_metric}_Mean: {mean_all_classes:.4f} ({class_details})")
168
+ epoch_metrics[f"val_{name_metric}"] = mean_all_classes
169
+ for idx, value in enumerate(avg_list):
170
+ epoch_metrics[f"val_{name_metric}"] = value
171
+
133
172
  else:
134
173
  avg_val = val_sum / num_batchs
135
174
  metrics_summary_str.append(f"{name_metric}: {avg_val:.4f}")
175
+ epoch_metrics[f"val_{name_metric}"] = avg_val
176
+
136
177
 
137
178
  formatted_metrics = " | ".join(metrics_summary_str)
138
179
  tqdm.write(f"Epoch {current_epoch + 1}/{epochs} — Validation: {formatted_metrics}")
139
-
180
+
181
+ if self.use_mlflow:
182
+ mlflow.log_metrics(epoch_metrics, step=current_epoch + 1)
183
+
184
+ if self.write_result_in_CSV:
185
+ with open(self.csv_path, 'a', newline='', encoding='utf-8') as file:
186
+ writer = csv.DictWriter(file, fieldnames=csv_fields, extrasaction='ignore')
187
+ writer.writerow(epoch_metrics)
188
+
140
189
  if save_checkpoint:
141
190
  self.save_model(f"checkpoint_{self.name}.pt")
142
191
 
@@ -190,11 +239,11 @@ class BaseModel(nn.Module):
190
239
  num_batches = len(eval_loader)
191
240
  final_metrics = {}
192
241
 
193
- for name_metrics, val_sum in dict_metrics_sum.items():
242
+ for current_metrics, val_sum in dict_metrics_sum.items():
194
243
  if isinstance(val_sum, list):
195
- final_metrics[name_metric] = [v / num_batches for v in val_sum]
244
+ final_metrics[current_metrics] = [v / num_batches for v in val_sum]
196
245
  else:
197
- final_metrics[name_metric] = val_sum / num_batches
246
+ final_metrics[current_metrics] = val_sum / num_batches
198
247
 
199
248
  return final_metrics
200
249
 
@@ -0,0 +1,322 @@
1
+ import torch
2
+ from huggingface_hub import hf_hub_download
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ from .base_model import BaseModel
6
+ from .lib.convolutions import DoubleConv
7
+ import math
8
+
9
+ class EfficientSelfAttention(nn.Module):
10
+ def __init__(self, dim, num_heads = 8, reduction_ratio = 1, dropout_p = 0.1):
11
+ super().__init__()
12
+ self.dim = dim # = C in the article
13
+ self.num_heads = num_heads
14
+ self.head_dim = dim // num_heads
15
+ self.scale = 1.0 / math.sqrt(self.head_dim)
16
+ self.reduction_ratio = reduction_ratio
17
+
18
+ self.W_q = nn.Linear(dim, dim)
19
+ if reduction_ratio > 1:
20
+ self.sr = nn.Conv2d(dim, dim, kernel_size=reduction_ratio, stride=reduction_ratio)
21
+ self.norm = nn.LayerNorm(dim)
22
+ else:
23
+ self.sr = None
24
+ self.norm = None
25
+
26
+ self.W_k = nn.Linear(dim, dim)
27
+ self.W_v = nn.Linear(dim, dim)
28
+ self.proj = nn.Linear(dim, dim)
29
+
30
+ self.attn_drop = nn.Dropout(dropout_p)
31
+ self.proj_drop = nn.Dropout(dropout_p)
32
+
33
+ def forward(self, X, H, W):
34
+ B, N, C = X.shape # N = h * W
35
+ # (B, N, C) -> (B, N, heads, head_dim) -> (B, heads, N, head_dim)
36
+ q = self.W_q(X).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
37
+ if self.sr is not None:
38
+ x_spatial = X.permute(0, 2, 1). reshape(B, C, H, W)
39
+ x_spatial = self.sr(x_spatial)
40
+ x_spatial = x_spatial.flatten(2) # B, C, H*W
41
+ x_spatial = x_spatial.permute(0, 2, 1)
42
+ x_spatial = self.norm(x_spatial)
43
+
44
+ k = self.W_k(x_spatial).reshape(B, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
45
+ v = self.W_v(x_spatial).reshape(B, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
46
+ else:
47
+ # Classical Attention
48
+ k = self.W_k(X).reshape(B, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
49
+ v = self.W_v(X).reshape(B, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
50
+
51
+ attn_value = (q @ k.transpose(-2, -1)) * self.scale
52
+ attn_value = attn_value.softmax(dim=-1)
53
+ attn_value = self.attn_drop(attn_value)
54
+
55
+ # (B, heads, N, N_reduced) @ v: (B, heads, N_reduced, head_dim) -> (B, heads, N, head_dim)
56
+ X_out = attn_value @ v
57
+ X_out = X_out.permute(0, 2, 1, 3). reshape(B, N, C)
58
+ proj = self.proj(X_out)
59
+ proj = self.proj_drop(proj)
60
+ return proj
61
+
62
+ class OverlapPatchMerging(nn.Module):
63
+ def __init__(self, in_channels, out_channels, kernel_size=3, stride=2, padding=1):
64
+ super().__init__()
65
+ self.conv = nn.Conv2d(in_channels=in_channels,
66
+ out_channels=out_channels,
67
+ kernel_size=kernel_size,
68
+ stride=stride,
69
+ padding=padding
70
+ )
71
+ self.norm = nn.LayerNorm(out_channels)
72
+ def forward(self, X):
73
+ # B, C, H, W
74
+ X = self.conv(X)
75
+ B, C, H_new, W_new = X.shape
76
+
77
+ X = X.flatten(2). transpose(2, 1) # B, N, C
78
+ X = self.norm(X)
79
+ return X, H_new, W_new
80
+
81
+
82
+ class TransformerSegFormer(nn.Module):
83
+ def __init__(self, N, C_in, C_out, num_heads, reduction_ratio = 1, intermediate_channel = 256, H = 256, W = 256, patch_kernel = 3, stride = 2, padding = 1):
84
+ super().__init__()
85
+
86
+ self.efficient_self_attns = nn.ModuleList()
87
+ self.mix_FFNs = nn.ModuleList()
88
+ self.N = N
89
+
90
+ for i in range(N):
91
+ self.efficient_self_attns.append(EfficientSelfAttention(dim = C_out, num_heads= num_heads, reduction_ratio= reduction_ratio))
92
+ target_H = H // stride
93
+ target_W = W // stride
94
+ self.mix_FFNs.append(MixFFN(C_out, intermediate_channel))
95
+
96
+ self.overlap_patch_merging = OverlapPatchMerging(C_in, C_out,patch_kernel, stride = stride, padding = padding)
97
+
98
+ def forward(self, X):
99
+ X, H_new, W_new = self.overlap_patch_merging(X)
100
+
101
+ for idx in range(self.N):
102
+ X = self.efficient_self_attns[idx](X, H_new, W_new) + X
103
+ X = self.mix_FFNs[idx](X, H_new, W_new)
104
+
105
+ return X, H_new, W_new
106
+
107
+ class MixFFN(nn.Module):
108
+ def __init__(self, in_channels, intermediate_channels, dropout_p = 0.1):
109
+ super().__init__()
110
+ self.MLP1 = nn.Linear(in_features=in_channels, out_features=intermediate_channels)
111
+ self.gelu = nn.GELU()
112
+ self.conv33 = nn.Conv2d(in_channels=intermediate_channels, out_channels=intermediate_channels, kernel_size=3, padding=1, groups=intermediate_channels) # Note: Conv 3x3 depthwise dans l'original
113
+ self.MLP2 = nn.Linear(in_features=intermediate_channels, out_features=in_channels)
114
+ self.dropout = nn.Dropout(dropout_p)
115
+ def forward(self, X, H, W):
116
+ skipped_X = X
117
+ B, N, C = X.shape
118
+
119
+ X = self.MLP1(X)
120
+ X = self.gelu(X)
121
+ X = self.dropout(X)
122
+
123
+ # Spatial reshape dynamique avec H et W passés en paramètre
124
+ X = X.transpose(1, 2).view(B, -1, H, W)
125
+ X = self.conv33(X)
126
+ X = X.flatten(2).transpose(1, 2)
127
+
128
+ X = self.MLP2(X)
129
+ X = self.dropout(X)
130
+ return X + skipped_X
131
+
132
+
133
+
134
+
135
+ class MLPLayers(nn.Module):
136
+ def __init__(self, i, input_channels, embbeding_channels, dropout = 0.1, H = 256, W = 256):
137
+ super().__init__()
138
+ # Keep the same dimension but project on the number of channels:
139
+ self.linear_project = nn.Conv2d(in_channels= input_channels, out_channels= embbeding_channels, kernel_size=1)
140
+
141
+ self.dropout = nn.Dropout(dropout)
142
+ self.current_H = H // (2**(i + 2))
143
+ self.current_W = W // (2**(i + 2))
144
+ self.target_H = H // 4
145
+ self.target_W = W // 4
146
+
147
+ self.linear_spatial = nn.Linear((self.current_H * self.current_W), self.target_H * self.target_W)
148
+
149
+ def forward(self, X):
150
+ # B, Ci, H/2**(i+1), W/2**(i+1)
151
+ B, _, _, _ = X.shape
152
+
153
+ # Projection on C:
154
+ # B, C, H/2**(i+1), W/2**(i+1)
155
+ # C is embbeding channels:
156
+ X = self.linear_project(X)
157
+
158
+ # B, C, H/2**(i+1)* W/2**(i+1)
159
+ X = X.flatten(2)
160
+
161
+ # B, C, H/4* W/4
162
+ X = self.linear_spatial(X)
163
+
164
+ X = self.dropout(X)
165
+ # B, C, H/4, W/4
166
+ X = X.view(B, -1, self.target_H, self.target_W)
167
+ return X
168
+
169
+
170
+ class SegFormerTorch(nn.Module):
171
+ def __init__(self, depth = 4, input_channels=1, output_channels=1, initial_filters=64, dropout_p = 0.1, H=256, W= 256):
172
+ super().__init__()
173
+ self.depth = depth
174
+ self.channels = [32, 64, 160, 256]
175
+ # self.channels = [initial_filters*2**i for i in range(self.depth)]
176
+ self.num_heads = [1, 2, 5, 8]
177
+ reduction_ratios = [8, 4, 2, 1]
178
+
179
+ self.transformers = nn.ModuleList()
180
+ current_in_channels = input_channels
181
+
182
+ current_H, current_W = H, W
183
+
184
+ for i in range(self.depth):
185
+ patch_kernel = 7 if i == 0 else 3
186
+ patch_stride = 4 if i == 0 else 2
187
+ patch_padding = 3 if i == 0 else 1
188
+
189
+
190
+ self.transformers.append(TransformerSegFormer(
191
+ N=2, # Nombre de blocs Attention+MixFFN par niveau
192
+ C_in=current_in_channels,
193
+ C_out=self.channels[i],
194
+ num_heads=self.num_heads[i],
195
+ reduction_ratio=reduction_ratios[i],
196
+ intermediate_channel=self.channels[i] * 4,
197
+ H=current_H,
198
+ W=current_W,
199
+ patch_kernel=patch_kernel,
200
+ stride=patch_stride,
201
+ padding=patch_padding
202
+ )
203
+ )
204
+ current_in_channels = self.channels[i]
205
+ current_H //= patch_stride
206
+ current_W //= patch_stride
207
+ self.embedding_dim = 256
208
+
209
+ self.mlp_layers = nn.ModuleList()
210
+ for i in range(self.depth):
211
+ self.mlp_layers.append(
212
+ MLPLayers(i=i, input_channels=self.channels[i], embbeding_channels=self.embedding_dim, H=H, W=W)
213
+ )
214
+ self.linear_fuse = nn.Conv2d(self.embedding_dim * self.depth, self.embedding_dim, kernel_size=1)
215
+ self.linear_pred = nn.Conv2d(self.embedding_dim, output_channels, kernel_size=1)
216
+
217
+ self.dropout = nn.Dropout(dropout_p)
218
+
219
+ def forward(self, X):
220
+ orig_H, orig_W = X.shape[-2], X.shape[-1]
221
+ features = []
222
+ for idx in range(self.depth):
223
+ X, H_new, W_new = self.transformers[idx](X)
224
+ X_spatial = X.transpose(1, 2).view(X.shape[0], -1, H_new, W_new).contiguous()
225
+ features.append(X_spatial)
226
+ X = X_spatial
227
+
228
+ decoder_output = []
229
+ for idx in range(self.depth):
230
+ decoder_output.append(self.mlp_layers[idx](features[idx]))
231
+
232
+ out = torch.cat(decoder_output, dim = 1)
233
+ out = self.linear_fuse(out)
234
+ out = F.gelu(out)
235
+ out = self.dropout(out)
236
+ out = self.linear_pred(out)
237
+ # To get the same dimension as the input:
238
+ out = F.interpolate(out, size=(orig_H, orig_W), mode='bilinear', align_corners=False)
239
+
240
+ return out
241
+
242
+
243
+
244
+ class SegFormer(BaseModel):
245
+ def __init__(self, input_channels=1, output_channels=1, H = 256, W = 256,depth=4, initial_filters=64, **kwargs):
246
+ super().__init__(name="SegFormer", **kwargs)
247
+ self.input_channels = input_channels
248
+ self.output_channels = output_channels
249
+ self.depth = depth
250
+ self.initial_filters = initial_filters
251
+ self.H = H
252
+ self.W = W
253
+
254
+ def create_model(self):
255
+ self.model = SegFormerTorch(
256
+ input_channels=self.input_channels,
257
+ output_channels=self.output_channels,
258
+ depth=self.depth,
259
+ initial_filters=self.initial_filters,
260
+ H = self.H,
261
+ W = self.W
262
+ )
263
+ return self.model
264
+
265
+
266
+ def load_nvidia_weights_into_custom_model(self, model_name="nvidia/mit-b0"):
267
+ if self.model is None:
268
+ raise ValueError("Please create_model first!")
269
+
270
+ try:
271
+ weights_path = hf_hub_download(repo_id=model_name, filename="pytorch_model.bin")
272
+ official_dict = torch.load(weights_path, map_location="cpu")
273
+ except Exception:
274
+ from safetensors.torch import load_file
275
+ weights_path = hf_hub_download(repo_id=model_name, filename="model.safetensors")
276
+ official_dict = load_file(weights_path)
277
+
278
+ custom_dict = self.model.state_dict()
279
+ new_state_dict = {}
280
+
281
+ for stage_idx in range(4):
282
+ src_prefix = f"segformer.encoder.patch_embeddings.{stage_idx}"
283
+ dst_prefix = f"transformers.{stage_idx}.overlap_patch_merging"
284
+
285
+ if f"{src_prefix}.proj.weight" in official_dict:
286
+ new_state_dict[f"{dst_prefix}.conv.weight"] = official_dict[f"{src_prefix}.proj.weight"]
287
+ new_state_dict[f"{dst_prefix}.conv.bias"] = official_dict[f"{src_prefix}.proj.bias"]
288
+ new_state_dict[f"{dst_prefix}.norm.weight"] = official_dict[f"{src_prefix}.layer_norm.weight"]
289
+ new_state_dict[f"{dst_prefix}.norm.bias"] = official_dict[f"{src_prefix}.layer_norm.bias"]
290
+ for block_idx in range(2):
291
+ src_block = f"segformer.encoder.block.{stage_idx}.{block_idx}"
292
+
293
+ # Efficient Self-Attention
294
+ dst_attn = f"transformers.{stage_idx}.efficient_self_attns.{block_idx}"
295
+ new_state_dict[f"{dst_attn}.W_q.weight"] = official_dict[f"{src_block}.attention.self.query.weight"]
296
+ new_state_dict[f"{dst_attn}.W_q.bias"] = official_dict[f"{src_block}.attention.self.query.bias"]
297
+ new_state_dict[f"{dst_attn}.W_k.weight"] = official_dict[f"{src_block}.attention.self.key.weight"]
298
+ new_state_dict[f"{dst_attn}.W_k.bias"] = official_dict[f"{src_block}.attention.self.key.bias"]
299
+ new_state_dict[f"{dst_attn}.W_v.weight"] = official_dict[f"{src_block}.attention.self.value.weight"]
300
+ new_state_dict[f"{dst_attn}.W_v.bias"] = official_dict[f"{src_block}.attention.self.value.bias"]
301
+ new_state_dict[f"{dst_attn}.proj.weight"] = official_dict[f"{src_block}.attention.output.dense.weight"]
302
+ new_state_dict[f"{dst_attn}.proj.bias"] = official_dict[f"{src_block}.attention.output.dense.bias"]
303
+
304
+ # Spatial Reduction (SR)
305
+ if f"{src_block}.attention.self.sr.weight" in official_dict:
306
+ new_state_dict[f"{dst_attn}.sr.weight"] = official_dict[f"{src_block}.attention.self.sr.weight"]
307
+ new_state_dict[f"{dst_attn}.sr.bias"] = official_dict[f"{src_block}.attention.self.sr.bias"]
308
+ new_state_dict[f"{dst_attn}.norm.weight"] = official_dict[f"{src_block}.attention.self.layer_norm.weight"]
309
+ new_state_dict[f"{dst_attn}.norm.bias"] = official_dict[f"{src_block}.attention.self.layer_norm.bias"]
310
+
311
+ # MixFFN
312
+ dst_ffn = f"transformers.{stage_idx}.mix_FFNs.{block_idx}"
313
+ new_state_dict[f"{dst_ffn}.MLP1.weight"] = official_dict[f"{src_block}.mlp.dense1.weight"]
314
+ new_state_dict[f"{dst_ffn}.MLP1.bias"] = official_dict[f"{src_block}.mlp.dense1.bias"]
315
+ new_state_dict[f"{dst_ffn}.conv33.weight"] = official_dict[f"{src_block}.mlp.dwconv.dwconv.weight"]
316
+ new_state_dict[f"{dst_ffn}.conv33.bias"] = official_dict[f"{src_block}.mlp.dwconv.dwconv.bias"]
317
+ new_state_dict[f"{dst_ffn}.MLP2.weight"] = official_dict[f"{src_block}.mlp.dense2.weight"]
318
+ new_state_dict[f"{dst_ffn}.MLP2.bias"] = official_dict[f"{src_block}.mlp.dense2.bias"]
319
+
320
+ missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
321
+ print(f"Loading successful ! Missing keys : {len(missing)}, Weird keys : {len(unexpected)}")
322
+ return self.model
@@ -7,7 +7,7 @@ from .lib.convolutions import DoubleConv
7
7
  class UnetPPTorch(nn.Module):
8
8
  def __init__(self, input_channels=1, output_channels=1, depth=4, initial_filters=64,
9
9
  kernel_size=3, padding=1):
10
- super(UnetPPTorch, self).__init__()
10
+ super(UnetPPTorch, self).__init__(name= "Unet++")
11
11
  self.depth = depth
12
12
  self.downs = nn.ModuleList()
13
13
  self.inter_down = nn.ModuleList()
@@ -1,43 +1,20 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: openmlcore
3
- Version: 1.0.0
3
+ Version: 1.2.0
4
4
  Summary: OpenMLCore is an open source framework (wrapper which simplify the use of PyTorch)
5
5
  Author-email: Doyez Sébastien <doyez.sebastien34090@gmail.com>
6
- License: MIT License
7
-
8
- Copyright (c) 2026 Votre Nom
9
-
10
- Permission is hereby granted, free of charge, to any person obtaining a copy
11
- of this software and associated documentation files (the "Software"), to deal
12
- in the Software without restriction, including without limitation the rights
13
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
- copies of the Software, and to permit persons to whom the Software is
15
- furnished to do so, subject to the following conditions:
16
-
17
- The above copyright notice and this permission notice shall be included in all
18
- copies or substantial portions of the Software.
19
-
20
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
- SOFTWARE.
27
6
  Project-URL: Homepage, https://sebastien-doyez2812.github.io/openMLCore.github.io/
28
7
  Classifier: Programming Language :: Python :: 3
29
8
  Classifier: License :: OSI Approved :: MIT License
30
9
  Classifier: Operating System :: OS Independent
31
10
  Requires-Python: >=3.10.0
32
11
  Description-Content-Type: text/markdown
33
- License-File: LICENSE
34
12
  Requires-Dist: torch
35
13
  Requires-Dist: tqdm
36
14
  Requires-Dist: onnx
37
15
  Requires-Dist: onnxscript
38
16
  Requires-Dist: torchinfo
39
17
  Requires-Dist: datasets
40
- Dynamic: license-file
41
18
 
42
19
  # OpenML-core, Open source Package for Machine learning
43
20
 
@@ -1,4 +1,3 @@
1
- LICENSE
2
1
  README.md
3
2
  pyproject.toml
4
3
  src/openmlcore/__init__.py
@@ -10,9 +9,11 @@ src/openmlcore.egg-info/top_level.txt
10
9
  src/openmlcore/base/data.py
11
10
  src/openmlcore/loss/loss.py
12
11
  src/openmlcore/metrics/metrics.py
12
+ src/openmlcore/models/MobileUnetR.py
13
13
  src/openmlcore/models/UNets.py
14
14
  src/openmlcore/models/attUnets.py
15
15
  src/openmlcore/models/base_model.py
16
+ src/openmlcore/models/segformer.py
16
17
  src/openmlcore/models/unetpp.py
17
18
  src/openmlcore/models/lib/attention.py
18
19
  src/openmlcore/models/lib/convolutions.py
@@ -1,11 +1,17 @@
1
+ import sys
1
2
  import pytest
2
3
  import torch
3
4
  import torch.nn as nn
4
5
  from torch.utils.data import DataLoader
5
6
  from torchgen import model
6
7
  from torchgen import model
7
- from openmlcore.models.base_model import *
8
- from openmlcore.metrics.metrics import *
8
+
9
+ from pathlib import Path
10
+ root_dir = Path(__file__).resolve().parent.parent
11
+ sys.path.append(str(root_dir))
12
+
13
+ from src.openmlcore.models.base_model import *
14
+ from src.openmlcore.metrics.metrics import *
9
15
 
10
16
  # class DummyModel(BaseModel):
11
17
  # def __init__(self):
@@ -0,0 +1,26 @@
1
+ import sys
2
+ from pathlib import Path
3
+ root_dir = Path(__file__).resolve().parent.parent
4
+ sys.path.append(str(root_dir))
5
+
6
+ from src.openmlcore.loss.loss import *
7
+
8
+ def test_BCEWithLogits_loss():
9
+ loss = BCEWithLogitsLoss(weights = [1.0])
10
+ assert loss is not None
11
+
12
+ def test_mean_squared_error_loss():
13
+ loss = MeanSquaredErrorLoss(weights = [1.0])
14
+ assert loss is not None
15
+
16
+ def test_dice_loss():
17
+ loss = DiceLoss(weights = [1.0])
18
+ assert loss is not None
19
+
20
+ def test_focal_loss():
21
+ loss = FocalLoss(weights = [1.0])
22
+ assert loss is not None
23
+
24
+ def test_huber_loss():
25
+ loss = HuberLoss(weights = [1.0])
26
+ assert loss is not None
openmlcore-1.0.0/LICENSE DELETED
@@ -1,21 +0,0 @@
1
- MIT License
2
-
3
- Copyright (c) 2026 Votre Nom
4
-
5
- Permission is hereby granted, free of charge, to any person obtaining a copy
6
- of this software and associated documentation files (the "Software"), to deal
7
- in the Software without restriction, including without limitation the rights
8
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
- copies of the Software, and to permit persons to whom the Software is
10
- furnished to do so, subject to the following conditions:
11
-
12
- The above copyright notice and this permission notice shall be included in all
13
- copies or substantial portions of the Software.
14
-
15
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
- SOFTWARE.
@@ -1,27 +0,0 @@
1
- from openmlcore.loss.loss import *
2
-
3
- def test_BCEWithLogits_loss():
4
- loss = BCEWithLogitsLoss()
5
- assert loss is not None
6
-
7
- def test_mean_squared_error_loss():
8
- loss = MeanSquaredErrorLoss()
9
- assert loss is not None
10
-
11
- def test_dice_loss():
12
- loss = DiceLoss()
13
- assert loss is not None
14
-
15
- def test_focal_loss():
16
- loss = FocalLoss()
17
- assert loss is not None
18
-
19
- def test_huber_loss():
20
- loss = HuberLoss()
21
- assert loss is not None
22
-
23
- def test_kl_divergence_loss():
24
- loss = KLDivergenceLoss()
25
- assert loss is not None
26
-
27
-
File without changes
File without changes