Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
# Split data into train-test data sets
X = df.loc[:,'image_names']
y = df.loc[:,'emergency_or_not']
# Split
train_x, val_x, train_y, val_y = train_test_split(X, y,
test_size = 0.1,
random_state = 27,
stratify=y)
# Train df
df_train = pd.DataFrame(columns=['image_name','category'])
df_train['image_name'] = train_x
df_train['category'] = train_y
# Test df
df_test= pd.DataFrame(columns=['image_name','category'])
df_test['image_name'] = val_x
df_test['category'] = val_y
df_train.reset_index(drop=True, inplace=True)
df_test.reset_index(drop=True, inplace=True)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment