Reimagining Image Segmentation using Active Contour: From Chan Vese Algorithm into a Proposal Novel Functional Loss Framework

TL;DR

This paper reimagines image segmentation by extending the Chan-Vese algorithm into a novel functional loss framework using active contours, with empirical validation on standard datasets and implementation in PyTorch.

Contribution

It introduces a new functional segmentation loss based on active contours, integrating the Chan-Vese algorithm into modern deep learning frameworks.

Findings

The proposed loss improves segmentation accuracy on benchmark datasets.

Comparison shows advantages over traditional loss functions.

Implementation details enable reproducibility and further research.

Abstract

In this paper, we present a comprehensive study and analysis of the Chan-Vese algorithm for image segmentation. We employ a discretized scheme derived from the empirical study of the Chan-Vese model's functional energy and its partial differential equation based on its level set function. We provide a proof of the results and an implementation using MATLAB. Leveraging modern computer vision methodologies, we propose a functional segmentation loss based on active contours, utilizing pytorch.nn.ModuleLoss and a level set based on the Chan-Vese algorithm. We compare our results with common computer vision segmentation datasets and evaluate the performance of classical loss functions against our proposed method. All code and materials used are available at https://github.com/gguzzy/chan_vese_functional_loss.

Figures

Tables

Table 1. Table 1: Filtered images: MSE and RMSE

Method	MSE	RMSE
Gaussian Filtered	0.029759	0.17251
Median Filtered	0.00061407	0.02478
Gaussian S&P Filtered	0.0030896	0.055584
Median S&P Filtered	0.0030237	0.054988
Median Gaussian S&P Filtered	0.0030237	0.054988

Equations

$$I_{\text{noisy}}=I+\mathcal{N}(\mu ,{\sigma }^2),$$

$$G(x,y)=\frac{1}{2\pi {\sigma }^2}{e}^{-\frac{x^2+y^2}{2{\sigma }^2}},$$

$$E(u,C)=$$

$$+\mu {\int }_{\Omega \setminus C}(u-I{)}^{2}dx$$

$$+\nu \text{Length}(C),$$

$$F(c_1,c_2,C)=$$

$$+\nu \text{Area}(\text{int}(C))$$

$$+{\lambda }_1{\int }_{\text{int}(C)}|u_0-c_1{|}^{2}dx$$

$$+{\lambda }_2{\int }_{\text{ext}(C)}|u_0-c_2{|}^{2}dx$$

$$|f(x)-f(y)|\le L\Vert x-y\Vert ,$$

$${\varphi }_0(x,y)=\sqrt{(x-x_0{)}^2+(y-y_0{)}^2}-r,$$

$${\varphi }_0(x,y)=\sqrt{{\left(\frac{x-x_0}{a}\right)}^2+{\left(\frac{y-y_0}{b}\right)}^2}-1,$$

$$\{\begin{array}{l}C=\{(x,y)\in \Omega :\Phi (x,y)=0\}\\ \text{inside}(C)=\{(x,y)\in \Omega :\Phi (x,y)>0\}\\ \text{outside}(C)=\Omega \backslash \text{inside}(C)=\{(x,y)\in \Omega :\Phi (x,y)<0\}\end{array}$$

$$F(c_1,c_2,\Phi )=$$

$$+\nu {\int }_{\Omega }H(\Phi (x,y))dxdy$$

$$+{\lambda }_1{\int }_{\Omega }|u_0(x,y)-c_1{|}^{2}H(\Phi (x,y))dxdy$$

$$+{\lambda }_2{\int }_{\Omega }|u_0(x,y)-c_2{|}^2(1-H(\Phi (x,y)))dxdy$$

$$\displaystyle\frac{\partial\Phi}{\partial t}=\delta(\Phi)\bigg{(}$$

$$\displaystyle-\lambda_{1}(u_{0}-c_{1})^{2}+\lambda_{2}(u_{0}-c_{2})^{2}\bigg{)}$$

$$\kappa (\Phi )=\frac{{\Phi }_{xx}{\Phi }_y^2-2{\Phi }_{xy}{\Phi }_x{\Phi }_y+{\Phi }_{yy}{\Phi }_x^2}{({\Phi }_x^2+{\Phi }_y^2{)}^{3/2}}$$

$$H_{ϵ}(x)$$

$${\delta }_{ϵ}(x)$$

$${\Delta }_x{\Phi }_{i,j}^n$$

$${\Delta }_y{\Phi }_{i,j}^n$$

$$\frac{{\Phi }_{i,j}^{n+1}-{\Phi }_{i,j}^n}{\Delta t}=$$

$$\displaystyle+\frac{\Delta_{y}^{+}\Phi^{n}_{i,j}+\Delta_{y}^{-}\Phi^{n}_{i,j}}{2h^{2}}\Bigg{)}$$

$$\displaystyle-\delta_{h}(\Phi^{n}_{i,j})\Big{(}v\lambda+\frac{1}{u_{c}^{i,j}}-\frac{1}{2\Phi^{n}_{i,j}-\lambda_{u_{c}^{i,j}}}\Big{)}\Bigg{]}.$$

$$C_1$$

$$C_2$$

$$C_3$$

$$C_4$$

$${\Phi }_{i,j}^{n+1}={\Phi }_{i,j}^n+\Delta t{\delta }_h({\Phi }_{i,j}^n)\mu \left(C_1+C_2+C_3+C_4\right).$$

$$Q=\sum \left(\left|{\Phi }_{i,j}^{n+1}-{\Phi }_{i,j}^n\right|<\Delta t\cdot h^2\right)$$

$${\Phi }_{i,j}^{n+1}={\Phi }_{i,j}^n-\Delta t\cdot \text{sign}(\Phi (x,y,t))\cdot G({\Phi }_{i,j}^n)$$

$$S_{i,j}=R_{i,j}\wedge G_{i,j}\wedge B_{i,j}$$

$$S_{i,j}=R_{i,j}\vee G_{i,j}\vee B_{i,j}$$

$$S_{i,j}=\text{Majority}(R_{i,j},G_{i,j},B_{i,j})$$

Full text

1 Problem overview

Image segmentation is the process of partitioning a digital image into multiple segments (sets of pixels). Hence, image segmentation consists into portioning a digital image into valuable and reliable pieces of information, easier to analyze. The whole picture does not contain only relevant data, in most computer vision scenarios, we are interested in a portion of area of pixel of the image, which could be represented as people or some certain objects we are interested in, as our cat who crossed the perimeter in the garden. Nowadays the aims are immeasurable.

Although it is a common problem, image segmentation has a wide set of resolution methods which aim at collecting set of pixels or a set of relevant pixels called objects.

These methods are in general (i.e. not limited to):

• •

Thresholding: This method involves partitioning an image based on pixel intensity. Pixels with intensity above a certain threshold are classified differently from those below, allowing for straightforward segmentation of images into foreground and background.

• •

Clustering: This technique partitions image pixels into clusters based on attributes such as color and intensity. K-means clustering is a notable example, categorizing pixels into clusters with similar values.

• •

Contour-based methods: These methods aim to find contours within an image, identifying boundaries between different image regions. The active contour model, or snakes, dynamically adjusts contours to outline object edges.

• •

Region-based methods: These methods segment an image based on the homogeneity of pixel regions. Techniques like Region Growing start from seed points and aggregate neighboring pixels with similar properties to form a segment.

In this proposed work, we used an active contour modeling using the Chan-Vese method, an active contour modeling method renowned for its efficacy in capturing complex object boundaries and independent on image intensities.

2 Proposed approach

2.1 Preprocessing

In our experiment, an image goes through a well-defined process of image preprocessing, moreover the experiments are run onto different scenarios to prove its effectiveness, hence the starting image is segmented using the proposed method and compared to the noisy images processed through the next proposed approeaches. Indeed noise is introduced in order to add a non linear factor to the problem, as in real world applications. Each image has been treated, introducing noise and removing those factors later, using the following techniques:

2.2 Data Augmentation: noise introduction to the dataset

Data augmentation plays a critical role in the development of machine learning models, particularly in the field of computer vision. By introducing variations in the training data, models can learn to generalize better and become more robust to unseen data. In this work, we focus on the incorporation of Gaussian noise and Salt-and-Pepper (S&P) noise as augmentation strategies to simulate real-world imperfections in images.

Gaussian Noise Addition: Gaussian noise, characterized by its normal distribution, is added to images to mimic the effect of random variations in pixel values. This is mathematically represented as:

[TABLE]

where $I$ is the original image, $\mathcal{N}(\mu,\sigma^{2})$ denotes the Gaussian distribution with mean $\mu$ and variance $\sigma^{2}$, and $I_{\text{noisy}}$ is the resultant image. This technique aids in training models that are less sensitive to slight pixel intensity variations.

Salt-and-Pepper Noise Injection: S&P noise, also known as impulse noise, introduces sharp, sudden disturbances in the image, represented by randomly scattering white and black pixels. This type of noise challenges the model’s ability to maintain performance despite the presence of significant pixel-level anomalies.

All these techniques aim at improving the model’s robustness, meaning the efficiency of the segmentation to low quality datasets.

2.3 Feature Engineering: Advanced Noise Removal Techniques

Effective feature engineering is paramount in preparing data to increase its data quality for our model, especially when dealing with noisy datasets. To counteract the noise introduced during data augmentation, we employ sophisticated noise removal techniques as part of our preprocessing pipeline.

• •

Gaussian Noise Elimination via Gaussian Filtering: Gaussian filtering is a smoothing technique that reduces Gaussian noise. It utilizes a Gaussian kernel to blur the image, effectively diminishing the impact of noise:

[TABLE]

where $G(x,y)$ is the Gaussian kernel, and $\sigma$ is the standard deviation of the distribution. This method preserves edge properties while reducing noise, making it ideal for preprocessing in image segmentation tasks.

• •

Salt-and-Pepper Noise Reduction via Median Filtering: Median filtering is a non-linear process used to remove S&P noise. It replaces each pixel’s value with the median value of the intensities in its neighborhood. This approach effectively preserves edges while removing noise, essential for maintaining the structural integrity of objects in the image.

As evidenced in our final benchmark tests in Figure 18, even with the introduction of random noise, the models maintain high performance levels. This underscores the effectiveness of our proposed augmentation and preprocessing strategies, demonstrating the model’s robustness against various noise types.

From Mumford-Shah to Chan-Vese

Introduced by Mumford and Shah in their seminal 1989 paper [1], the Mumford-Shah model has been a cornerstone in the field of image segmentation. The model seeks to partition an image into regions of approximately uniform intensity by minimizing an energy functional. For a given image $I$ defined over a domain $\Omega\in\mathbb{R}^{2}$, the Mumford-Shah functional can be expressed as:

The Mumford-Shah model

[TABLE]

where:

• •

$u$ is an approximation of the image

• •

$I$ within the domain $\Omega\setminus C$

• •

$C$ denotes the set of edges in the image

• •

$\mu$, $\nu$ are regularization parameters, chosen arbitrarily by the user.

In the Mumford-Shah functional, $\Omega$ represents the entire image domain, while $\Omega\setminus C$ denotes the image domain excluding the set of edges $C$, essentially covering the regions within the image that are not part of the segmentation boundaries. The parameter $\mu$ controls the trade-off between the image’s fidelity to its approximation $u$ and smoothness, whereas $\nu$ regulates the importance of the contour’s length in the overall segmentation, aiming to minimize unnecessary segmentation complexity. The functional aims to balance the smoothness of $u$ and the fidelity to the original image $I$, along with the complexity of the segmentation represented by $\text{Length}(C)$.

The model aims to find piecewise smooth representations of images by minimizing an energy function that penalizes the image’s deviation from the model and the length of the contours.

Expanding upon the Mumford-Shah model, Chan and Vese introduced a simplified version in 2001 [2], targeting images with two distinct intensity levels. This model, known as the Chan-Vese model, is designed for binary segmentation tasks, optimizing a similar energy functional to directly distinguish between the two regions [2].

While the original Mumford-Shah model does not inherently rely on level set functions, its adaptations, including the Chan-Vese model, often utilize level set methods to represent the evolving curves in the segmentation process. This approach allows for a more flexible representation of contours and simplifies the numerical implementation of the segmentation task.

The Chan-Vese model

The Chan-Vese model is a specialization of the Mumford-Shah model for segmenting images into two distinct regions. It uses an energy function defined as:

[TABLE]

where $c_{1}$ and $c_{2}$ represent the mean intensities inside and outside the contour $C$. The terms $\lambda_{1}$, $\lambda_{2}$, and $\mu$ are weighting parameters that control the segmentation’s adherence to the model’s assumptions.

This model effectively segments images by evolving a contour $C$ to minimize the above energy, adapting well to images with homogenous regions.

The Chan-Vese model simplifies this approach by assuming images with two intensity levels, optimizing a similar energy function to distinguish between these regions directly.

Example: The Chan-Vese algorithm excels in segmenting images with clear object boundaries against a uniform background, efficiently handling variations in intensity and texture without prior knowledge of object location.

2.4 Level Set Methods

Level set methods play a pivotal role in the domain of computational image segmentation, offering a robust framework for dynamically evolving contours. Within the Chan-Vese model, these methods facilitate the optimization of an energy functional that guides the segmentation process. The level set function, $\phi(x,y,t)$, encapsulates the contour $C$ implicitly as its zero level set, evolving under a carefully designed partial differential equation (PDE). This evolution is meticulously driven to minimize the energy functional, as detailed in Equation 4.

where $c_{1}$ and $c_{2}$ denote the mean intensities within and outside the evolving contour $C$, respectively. The weighting parameters $\lambda_{1}$, $\lambda_{2}$, and $\mu$ are calibrated to balance the trade-offs inherent in the segmentation task, ensuring an adherence to the image’s inherent structures while promoting contour smoothness and minimal complexity.

2.4.1 The Lipschitz Function

A critical aspect of the level set formulation is the initialization of the level set function $\phi_{0}(x,y)$, which should ideally be a Lipschitz function. This property is crucial for ensuring the numerical stability and convergence of the evolving contour. A Lipschitz function adheres to the condition:

[TABLE]

for every pair $x,y$ within the domain, where $L$ represents the Lipschitz constant. This constraint ensures a controlled evolution pace, preventing excessive, abrupt changes in the contour’s geometry.

Initialization with a Circle or Ellipse: For initializing the level set function in the Chan-Vese framework, employing geometric shapes such as circles or ellipses offers a pragmatic and theoretically sound approach. Both configurations qualify as Lipschitz functions, attributable to their smooth, continuous boundaries that can be mathematically articulated with bounded derivatives.

A circular initial contour centered at $(x_{0},y_{0})$ with radius $r$ is defined by:

[TABLE]

demonstrating Lipschitz continuity as the magnitude of its gradient is uniformly bounded [5]. Similarly, an elliptical contour, characterized by its semi-major and semi-minor axes $a$ and $b$, respectively, the subtraction by 1 adjusts the level set to position the zero level set contour exactly at the boundary of the ellipse. This is the correct formulation for representing an elliptical contour in terms of a level set function; it is represented as:

[TABLE]

which also exhibits bounded gradients, affirming its suitability as an initial Lipschitz condition for the level set evolution.

The employment of these geometric forms as initial conditions not only simplifies the numerical implementation but also aligns with the theoretical underpinnings required for a stable and efficient segmentation process. This approach underscores the Chan-Vese model’s versatility in adapting to images with diverse intensity profiles, facilitating precise segmentation outcomes.

Level Set Formulation

Instead of searching for the solution in terms of $C$, we can redefine the problem in the level set formalism. In the level set method, $C\subseteq\Omega$ is represented by the zero level set of some Lipschitz function $\Phi:\Omega\rightarrow\mathbb{R}$, such that:

[TABLE]

Energy Functional

The energy functional in terms of $\Phi$ is given by:

[TABLE]

where:

• •

$\mu$ is the parameter that weights the length of the contour, controlling the smoothness of the segmentation boundary;

• •

$\nu$ weights the area inside the contour, affecting how the area influences the segmentation;

• •

$\lambda_{1}$ and $\lambda_{2}$ weight the intensity differences within and outside the contour, respectively, aiding in fitting the model to the image data;

• •

$H$ is the Heaviside function, used to differentiate the regions inside and outside the contour;

• •

$x$ and $y$ are the spatial coordinates in the image domain $\Omega$;

• •

$\nabla\Phi$ is the gradient of $\Phi$, with its norm representing the rate of change of $\Phi$, related to the curvature of the contour $C$ in the level set method.

• •

The domain $\Omega$ represents the area over which the image is defined and the segmentation is performed.

Partial Differential Equation for $\Phi$

The function $\Phi$ which minimizes the energy functional satisfies the PDE:

[TABLE]

The curvature $\kappa(\Phi)$ [3] of the evolving contour is given by the spatial derivatives of $\Phi$ up to the second order:

[TABLE]

Numerical Scheme

First, we introduce the regularization of the Heaviside and Dirac delta functions for some constant $\epsilon>0$ as follows:

[TABLE]

These regularizations are used in simulations to lead to the global minimum of the energy efficiently. Choosing $\epsilon=h$, where $h$ is the space step, is reasonable since $h$ is the smallest space step in the problem. In the available case we used $\epsilon=h=1$.

Discretization of the PDE

Let $\Phi^{n}_{i,j}=\Phi(n\Delta t,x_{i},y_{j})$, where $\Delta t$ is the time step. The PDE can be discretized using spatial finite differences as follows:

[TABLE]

The linearized discretized PDE becomes:

[TABLE]

Constants Definition

Defining constants for a given $\Phi^{n}$ as:

[TABLE]

The simplified equation then is:

[TABLE]

Summary of the Algorithm

Initialize $\Phi^{0}_{i,j}$ to some Lipschitz function $\Phi_{0}$. 2. 2.

Compute $c_{1}$ and $c_{2}$ using the regularized Heaviside function. 3. 3.

Solve the PDE using the discretized equation. 4. 4.

Reinitialize $\Phi^{n+1}_{i,j}$ to be the signed distance function to its zero level set using the reinitialization equation. 5. 5.

Check whether the solution is stationary. If not, continue; otherwise, stop.

The process of evolving the level set function $\Phi$ during image segmentation aims to minimize an energy functional that accurately represents the object’s boundary within the image. This iterative update of $\Phi$ continues until it converges to a stationary state, which signifies that the contour of the object has been precisely delineated.

The update mechanism for $\Phi$ at each grid point $(i,j)$ from iteration $n$ to $n+1$ relies on the image data and the curvature of $\Phi$ itself. This iterative update is terminated when the changes in $\Phi$ between successive iterations become insignificant across the entire image domain, implying that $\Phi$ has reached a stationary state and further iterations would not meaningfully alter its configuration. The mathematical expression for this stopping criteria is as follows:

The iterative process ceases when the discrepancy between $\Phi^{n+1}_{i,j}$ and $\Phi^{n}_{i,j}$ is lesser than $\Delta t\cdot h^{2}$, indicating the attainment of a stationary state. This condition is verified through the computation of a quantity $Q$, which assesses the stationarity of $\Phi$ across iterations:

[TABLE]

In this context, $\Delta t$ symbolizes the timestep size, and $h$ signifies the spatial resolution of the grid. The summation traverses all grid points $(i,j)$, with the inner comparison producing a Boolean value that denotes whether the change at each point is beneath the predefined threshold. The algorithm concludes when a significant majority of the points meet this stipulation, suggesting that $\Phi$ has become stationary.

For the purpose of reinitializing $\Phi$, the ensuing equation is utilized to guarantee the maintenance of the signed distance function property, which is pivotal for the numerical stability and precision of the level set method:

[TABLE]

Herein, $\text{sign}(\Phi(x,y,t))$ ensures the directionality of the update respects the gradient of the level set function, aiming to minimize deviations from a signed distance function. Concurrently, $G(\Phi^{n}_{i,j})$ is a function derived from the gradient of $\Phi$, adjusted to realign its values towards sustaining the signed distance property.

2.5 Chan-Vese Loss Implementation for RGB Images in PyTorch

In the pursuit of enhancing image segmentation models, we introduce the ChanVeseLossPyTorchRGB, a PyTorch module designed to adapt the classical Chan-Vese loss for RGB images. This module is particularly tailored for segmentation tasks, leveraging the characteristics of RGB images to compute the loss more effectively.

2.5.1 Implementation Details

The PyTorch module ChanVeseLossPyTorchRGB implements this adapted loss function. It requires the predicted segmentation masks, the ground truth masks, and the original RGB images as inputs. The intensity difference loss, $\mathcal{L}_{\text{intensity},c}$, is computed for each RGB channel and averaged. The smoothness term, $\mathcal{L}_{\text{smoothness}}$, utilizes the gradient of the predicted mask, emphasizing the regularization of the segmentation boundary.

The key components of this module are:

• •

The heaviside function, approximated using the hyperbolic tangent to ensure a smooth transition, facilitating the differentiation process required for gradient descent optimization.

• •

The mean_intensities function, which calculates the average intensities inside and outside the segmented region for each channel, crucial for the intensity difference computation.

• •

The computation of the smoothness loss, employing the gradient of the sigmoid-activated predictions, to enforce boundary smoothness.

The incorporation of this loss function into segmentation models promises enhanced performance by effectively leveraging the color information inherent in RGB images and enforcing smooth, coherent segmentations.

Python code snippet for the ChanVeseLossPyTorchRGB module:

See code here.

3 Results

A solution in MATLAB code is provided at the end for initialization of level set functions, the evolution of the curvature using the above discretized method, able to cope the dinamic solution of numerical systems.

Consider some assumptions have been done, as for instance using a circle as initial level set method for final test, since it is common use in the literature (look for reference and citation), moreovfer $dy=dx=1$ has been consider to 1 supposing that images are equally spaced, $epsilon=h=1$ since we are dealing with square images, each step should be of length 1 (moving of one step at the time). At last, the area is set to 0 since it doesnt efficiently help with the image segmentation resolution.

Here below there are few experimental results from the aforementioned resolution method:

3.1 Grey images

3.2 RGB images

When processing RGB images, as in LABEL:fig:rgbImages with the segmentation algorithm, it is imperative to consider the distinct nature of each RGB channel. Consequently, the algorithm is executed separately on each of the RGB channels, resulting in three preliminary segmentation maps. To synthesize these into a final, comprehensive segmentation result, an aggregation step is employed, utilizing logical operations such as AND, OR, and Majority Vote. The aggregation process is delineated as follows:

• •

Logical AND: This operation aggregates the segmentation maps by considering a pixel in the final segmentation to be part of the object if and only if it is identified as part of the object in all three RGB channel segmentations. Mathematically, for a given pixel position $(i,j)$, the final segmentation $S_{i,j}$ is defined as:

[TABLE]

where $R_{i,j}$, $G_{i,j}$, and $B_{i,j}$ represent the segmentation decisions at pixel $(i,j)$ for the Red, Green, and Blue channels, respectively.

• •

Logical OR: In contrast, the OR operation aggregates the maps by marking a pixel as part of the object if it is identified as such in any one of the RGB channel segmentations. Thus, for each pixel $(i,j)$:

[TABLE]

• •

Majority Vote: This approach determines a pixel’s inclusion in the final segmentation based on a majority rule among the three channels. A pixel is considered part of the object if at least two out of the three channel segmentations agree on its inclusion. Formally:

[TABLE]

where $\text{Majority}(R_{i,j},G_{i,j},B_{i,j})$ returns the majority value among $R_{i,j}$, $G_{i,j}$, and $B_{i,j}$.

Following the segmentation of each RGB image channel and the subsequent aggregation via these logical operations, the resulting segmentation maps are illustrated below each corresponding RGB image, showcasing the impact of the chosen aggregation method on the final segmentation outcome.

Here below an example of generated mask.

4 Appendix A: Project results visualization

4.1 Preprocessing and data preparation

As we can establish from Section 2. ’Proposed Approach’, within the Preprocessing and Data Augmentation part, we can see how the image data is treated. Recall that a digital image is stored as $I\in\mathbf{R}^{r\times c\times n_{channels}}$, where $r$ represents the image length, whereas $c$ (# columns), and $N_{channels}$ is the number of channels of colored images, in this case since we dealt with RGB images, it is indeed $I\in\mathbf{R}^{128\times 128\times 3}$. Recall that for grey scaled images, the task is a bit easier, since the level method evolution works only on a unique channel, instead of RGB images where multiple tasks has to run for each channel, and its result is made up by the aggregation of those 3 channels. Various techniques are exploited, as shown in subsection 3.2, here below you can find experimental results on the dataset.

5 Appendix B: Additional run experiments examples

Here below in figure 20 we can also see an additional result run using MRI (Magnetic resonance imaging) as initial input as shown in figure 21.

As we can see from the segmented mask, the area in black filled with unsual white patterns shows us a tumor alike pattern.

6 Discussion

In this paper we saw how the Chan-Vese method can be applied to a wide range of images accomplishing overall good results, even with images subject to various types of noise. The project could be expanded by exploiting the loss implemented both in MATLAB and in python is possible to train on a custom dataset to further explore better optimization techniques for increased performances in the loss module, in contrast to new SoAT methodologies.