{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os, ast\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nfrom sklearn.model_selection import StratifiedKFold","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join('/kaggle/input/global-wheat-detection/', 'train.csv'))\n# Let's expand the bounding box coordinates and calculate the area of all the bboxes\ntrain_df[['x_min','y_min', 'width', 'height']] = pd.DataFrame([ast.literal_eval(x) for x in train_df.bbox.tolist()], index= train_df.index)\ntrain_df = train_df[['image_id', 'bbox', 'source', 'x_min', 'y_min', 'width', 'height']]\ntrain_df['area'] = (train_df['width'] * train_df['height']).astype(np.int64)\ntrain_df['x_max'] = (train_df['x_min'] + train_df['width']).astype(np.int64)\ntrain_df['y_max'] = (train_df['y_min'] + train_df['height']).astype(np.int64)\ntrain_df = train_df.drop(['bbox'], axis=1)\ntrain_df[['x_min', 'y_min', 'width', 'height']] = train_df[['x_min', 'y_min', 'width', 'height']].astype(np.int64)\ntrain_df = train_df[['image_id', 'x_min', 'y_min', 'x_max', 'y_max', 'width', 'height', 'area', 'source']]\n\n# There are some buggy annonations in training images having huge bounding boxes. Let's remove those bboxes\ntrain_df = train_df[train_df['area'] < 100000]\n\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Add a new column to store kfold indices\ntrain_df.loc[:, 'kfold'] = -1\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_source = train_df[['image_id', 'source']].drop_duplicates()\n\n# get lists for image_ids and sources\nimage_ids = image_source['image_id'].to_numpy()\nsources = image_source['source'].to_numpy()\n\nskf = StratifiedKFold(n_splits=5, shuffle=True, random_state=1996)\nsplit = skf.split(image_ids, sources) # second arguement is what we are stratifying by\n\nfor fold, (train_idx, val_idx) in enumerate(split):\n    translated_val_idx = train_df[train_df['image_id'].isin(image_ids[val_idx])].index.values\n    print(len(translated_val_idx))\n    train_df.loc[translated_val_idx, 'kfold'] = fold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.to_csv('train.csv', index=False)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}