File size: 1,058 Bytes
c176aea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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
|