--- a
+++ b/src/hs_dataset.py
@@ -0,0 +1,58 @@
+"""
+Developed by: Daniel Crovo
+Dataset class definition
+
+"""
+from torch.utils.data import Dataset
+import os
+from PIL import Image 
+from pydicom import dcmread
+from dicom_utils import apply_windowing
+import numpy as np
+
+class HSDataset(Dataset): 
+    def __init__(self, image_dir, mask_dir, transform = None, w_level=35, w_width=350, normalized =False) -> None:
+        """Initialises the Heart Segmentation Dataset class. The dataset asumes all the dicom files are in one single folder aswell as the masks
+            both the images (dicom files) and masks should have the same name for each slice
+        Args:
+            image_dir (string): Path to the dicom directory
+            mask_dir (string): Path to the masks directory 
+            transform (_type_, optional):Aplied transformations Defaults to None.
+            w_level (int, optional): The window level (center). Defaults to 35.
+            w_width (int, optional): The width of the window Defaults to 400.
+            normalized (bool, optional): wheter to normalized the windowed array after applying windoing
+        """
+        super().__init__()
+        
+        self.image_dir = image_dir
+        self.mask_dir = mask_dir
+        self.transform = transform
+        self.w_level = w_level
+        self.w_width = w_width
+        self.normalized = normalized
+        self.images = os.listdir(self.image_dir)
+        self.masks = os.listdir(self.mask_dir)
+
+    def __len__(self):
+        return len(self.images)
+    
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.image_dir, self.images[idx])
+        mask_path = os.path.join(self.mask_dir, self.masks[idx])
+        ds = dcmread(img_path)
+        dicom_img = apply_windowing(ds = ds,window_center = self.w_level, 
+                                    window_width =self.w_width, normalized = self.normalized)
+        image = Image.fromarray(dicom_img).convert('RGB')
+        image = np.array(image)
+        mask = Image.open(mask_path).convert('L')
+        mask = np.array(mask)
+        mask[mask == 255.0] = 1.0
+
+
+        if self.transform is not None:
+            transformations = self.transform(image = image, mask = mask)
+            image = transformations['image']
+            mask = transformations['mask']
+            
+        return image, mask
+