| 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 | |