[lwnet] Add parameter for crop size

src/mednet/libs/segmentation/config/models/lwnet.py

@@ -23,4 +23,5 @@ model = LittleWNet(
     optimizer_type=Adam,
     optimizer_arguments=dict(lr=max_lr),
     augmentation_transforms=[],
+    crop_size=544,
 )

src/mednet/libs/segmentation/models/lwnet.py

@@ -260,6 +260,8 @@ class LittleWNet(pl.LightningModule):
         applied on the input **before** it is fed into the network.
     num_classes
         Number of outputs (classes) for this model.
+    crop_size
+        The size of the image after center cropping.
     """
 
     def __init__(
@@ -270,13 +272,14 @@ class LittleWNet(pl.LightningModule):
         optimizer_arguments: dict[str, typing.Any] = {},
         augmentation_transforms: TransformSequence = [],
         num_classes: int = 1,
+        crop_size: int = 544,
     ):
         super().__init__()
 
         self.name = "lwnet"
         self.num_classes = num_classes
 
-        self.model_transforms = [CenterCrop(size=(544, 544))]
+        self.model_transforms = [CenterCrop(size=(crop_size, crop_size))]
 
         self._train_loss = train_loss
         self._validation_loss = (