IN-P2/src/P2/preprocessing.py

61 lines
1.6 KiB
Python

from pandas import DataFrame, read_csv
from sklearn.preprocessing import normalize
def replace_values(df):
for column in df.columns:
df[column].fillna(value=df[column].mean(), inplace=True)
return df
def process_na(df, action):
if action == "drop":
return df.dropna()
elif action == "fill":
return replace_values(df)
else:
print("Unknown action selected. The choices are: ")
print("fill: fills the na values with the mean")
print("drop: drops the na values")
exit()
def filter_dataframe(df):
relevant_columns = [
"HORA",
"DIASEMANA",
"COMUNIDAD_AUTONOMA",
"ISLA",
"TOT_HERIDOS_LEVES",
"TOT_HERIDOS_GRAVES",
"TOT_VEHICULOS_IMPLICADOS",
"TOT_MUERTOS",
"TIPO_VIA",
"LUMINOSIDAD",
"FACTORES_ATMOSFERICOS",
]
filtered_df = df.filter(items=relevant_columns)
return filtered_df
def normalize_numerical_values(df):
cols = [
"TOT_HERIDOS_LEVES",
"TOT_HERIDOS_GRAVES",
"TOT_VEHICULOS_IMPLICADOS",
"TOT_MUERTOS",
]
filtered_df = df.filter(items=cols)
normalized_data = normalize(X=filtered_df)
normalized_df = DataFrame(data=normalized_data, columns=filtered_df.columns)
df.update(normalized_df)
return df
def parse_data(source, action):
df = read_csv(filepath_or_buffer=source, na_values="?")
processed_df = process_na(df=df, action=action)
filtered_df = filter_dataframe(df=processed_df)
normalized_df = normalize_numerical_values(df=filtered_df)
return normalized_df