ium_434766/create.py
2021-04-10 13:22:11 +02:00

51 lines
1.7 KiB
Python

import os
import numpy as np
import pandas as pd
import wget
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
def downloadCSV():
url = 'https://git.wmi.amu.edu.pl/s434766/ium_434766/raw/branch/master/healthcare-dataset-stroke-data.csv'
wget.download(url, out='healthcare-dataset-stroke-data.csv', bar=None)
def dropNaN():
data = pd.read_csv('healthcare-dataset-stroke-data.csv')
data = data.dropna()
return data
def NormalizeData(data):
data = data.astype({"age": np.int64})
for col in data.columns:
if data[col].dtype == object: # STRINGS TO LOWERCASE
data[col] = data[col].str.lower()
if data[col].dtype == np.float64: # FLOATS TO VALUES IN [ 0, 1]
dataReshaped = data[col].values.reshape(-1,1)
scaler = MinMaxScaler(feature_range=(0, 1))
data[col] = scaler.fit_transform(dataReshaped)
if col == 'ever_married': # YES/NO TO 1/0
data[col] = data[col].map(dict(yes=1, no=0))
if col == 'smoking_status':
data[col] = data[col].str.replace(" ", "_")
if col == 'work_type':
data[col] = data[col].str.replace("-", "_")
return data
def saveToCSV(data1,data2,data3):
data1.to_csv("data_train.csv", index=False)
data2.to_csv("data_test.csv",index=False)
data3.to_csv("data_val.csv",index=False)
downloadCSV()
data = dropNaN()
data = NormalizeData(data)
data_train, data_test = train_test_split(data, test_size=0.2, random_state=1)
data_train, data_val = train_test_split(data_train, test_size=0.25, random_state=1) ## Twice to get 0.6, 0.2, 0.2
saveToCSV(data_train,data_test,data_val)