code: def dataset(root_path):
      mask_files = glob.glob(root_path + "/*/*_mask*")
      image_files = [file.replace("_mask", "") for file in mask_files]
      files_df = pd.DataFrame(
          {
              "image_path": image_files,
              "mask_path": mask_files,
              "diagnosis": [diagnosis(x) for x in mask_files],
          }
      )
      train_df, val_df = train_test_split(
          files_df,
          stratify=files_df["diagnosis"],
          test_size=0.1,
          random_state=0,
      )
      train_df = train_df.reset_index(drop=True)
      val_df = val_df.reset_index(drop=True)
  
      train_df, test_df = train_test_split(
          train_df,
          stratify=train_df["diagnosis"],
          test_size=0.15,
          random_state=0,
      )
      train_df = train_df.reset_index(drop=True)
      test_df = test_df.reset_index(drop=True)
      print(
          "Train: {}\nVal: {}\nTest: {}".format(
              train_df.shape, val_df.shape, test_df.shape
          )
      )
  
      return train_df, val_df, test_df
  
imports: from __future__ import print_function
  from sklearn.model_selection import train_test_split
  import matplotlib.pyplot
  import albumentations
  import pandas
  import argparse
  import glob
  import os
  import torch
  import torch.nn
  import torch.optim
  import torch.utils.data
  from torchvision.utils import make_grid
  from bigdl.dllib.utils.log4Error import invalidInputError
  from bigdl.orca import init_orca_context, stop_orca_context
  from bigdl.orca.learn.pytorch import Estimator
  from Unet import UNet
  from dataset import *
label_desc: Use image node_id 1 for calling a global function with example usage: dataset(root_path) and returns: train_df,  val_df,  test_df