| from utils.data_augmentation import dataset |
| import os |
| import _pickle |
| import pandas as pd |
|
|
|
|
|
|
|
|
| def get_dataset(raw:bool=False, sample_size:int=1000, name:str='dataset.pkl',source:str='dataset.csv',boundary_conditions:list=None) -> _pickle: |
| """ Gets augmented dataset |
| |
| Args: |
| raw (bool, optional): either to use source data or augmented. Defaults to False. |
| sample_size (int, optional): sample size. Defaults to 1000. |
| name (str, optional): name of wanted dataset. Defaults to 'dataset.pkl'. |
| boundary_conditions (list,optional): y1,y2,x1,x2. |
| Returns: |
| _pickle: pickle buffer |
| """ |
| print(os.listdir('./data')) |
| if not(raw): |
| if name not in os.listdir('./data'): |
| ldat = dataset(sample_size,name,source,boundary_conditions) |
| ldat.generate() |
| with open(f"./data/{name}", "rb") as input_file: |
| buffer = _pickle.load(input_file) |
| else: |
| with open(f"./data/{source}", "rb") as input_file: |
| buffer = pd.read_csv(input_file) |
| return buffer |
|
|
|
|
