highD_data_process.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "import numpy as np\n",
    "from scipy.sparse import coo_matrix, csr_matrix\n",
    "from datatool import ade,vid_is_unique,foot2meter,vehicle2track,reset_idx,getNeighborGraph,graph2seq,get_displacement,train_test_val_split,matlab2dataframe\n",
    "from collections import Counter\n",
    "import time\n",
    "import pickle as pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_change_label(data):\n",
    "    data = data.sort_values(by=[\"id\",\"frame\"])\n",
    "    changes = []\n",
    "    for vid, df in data.groupby(\"id\"):\n",
    "        lane_id = df.laneId.to_numpy()\n",
    "        diff = lane_id[1:] - lane_id[:-1]\n",
    "        changed_or_not = (diff!=0)*1 # 1 : change; 0 : unchange\n",
    "        changed_or_not = np.hstack((np.zeros(1),changed_or_not))\n",
    "        changes.append(changed_or_not)\n",
    "    change_result = np.hstack(changes)\n",
    "    return change_result\n",
    "\n",
    "def getNeighborGraph(data,radius=10):\n",
    "    x,y,v_id,f_id = data.x,data.y,data.id,data.frame\n",
    "\n",
    "    vehicle_num, frame_num = v_id.max()+1, f_id.max()+1\n",
    "    sparse_X = csr_matrix((x, (v_id, f_id)), shape=(int(vehicle_num), int(frame_num))) # i行:车id;j列:时间;元素为i车j时刻的坐标x\n",
    "    sparse_Y = csr_matrix((y, (v_id, f_id)), shape=(int(vehicle_num), int(frame_num))) # i行:车id;j列:时间;元素为i车j时刻的坐标y\n",
    "    I_mat = (sparse_X!=0)*1 # i行:车id;j列:时间;元素为i车j时刻是否出现,出现为1,否则为0\n",
    "    mask = []\n",
    "    for v in range(I_mat.shape[0]):\n",
    "        concurrent_mask = I_mat.multiply(I_mat[v]) #同或 [1,0,1,0,0,0,1] & [1,0,1,1,1,0,0] = [1,0,1,0,0,0,0]\n",
    "\n",
    "        # 邻居xy坐标\n",
    "        concurrent_X = concurrent_mask.multiply(sparse_X) \n",
    "        concurrent_Y = concurrent_mask.multiply(sparse_Y)\n",
    "\n",
    "        # 自己xy坐标\n",
    "        self_x = concurrent_mask.multiply(sparse_X[v])\n",
    "        self_y = concurrent_mask.multiply(sparse_Y[v])\n",
    "\n",
    "        # 差值\n",
    "        delta_x = self_x - concurrent_X\n",
    "        delta_y = self_y - concurrent_Y\n",
    "\n",
    "        # 邻居x坐标在半径以内的指示矩阵\n",
    "        x_in_id = np.where((delta_x.data>-radius) & (delta_x.data<radius))\n",
    "        xc = delta_x.tocoo()\n",
    "        xrow_in = xc.row[x_in_id]\n",
    "        xcol_in = xc.col[x_in_id]\n",
    "        xI_data = np.ones(xrow_in.shape[0])\n",
    "        xneighbor_in_mat = csr_matrix((xI_data, (xrow_in, xcol_in)), shape=(I_mat.shape[0], I_mat.shape[1]))\n",
    "\n",
    "        # 邻居y坐标在半径以内的指示矩阵\n",
    "        y_in_id = np.where((delta_y.data>-radius) & (delta_y.data<radius))\n",
    "        yc = delta_y.tocoo()\n",
    "        yrow_in = yc.row[y_in_id]\n",
    "        ycol_in = yc.col[y_in_id]\n",
    "        yI_data = np.ones(yrow_in.shape[0])\n",
    "        yneighbor_in_mat = csr_matrix((yI_data, (yrow_in, ycol_in)), shape=(I_mat.shape[0], I_mat.shape[1]))\n",
    "\n",
    "        neighbor_in_mat = xneighbor_in_mat.multiply(yneighbor_in_mat).tolil()\n",
    "        neighbor_in_mat[v] = I_mat[v]\n",
    "        mask.append(neighbor_in_mat.tocsr())\n",
    "    return mask\n",
    "\n",
    "def graph2seq(data,graph_list,seq_length=16,max_vnum=30,down_sample_rate=5,sort_func=\"distance\"):\n",
    "    x,y,v_id,f_id,l = data.x,data.y,data.id,data.frame,data.label\n",
    "    vehicle_num, frame_num = v_id.max()+1, f_id.max()+1\n",
    "    sparse_X = csr_matrix((x, (v_id, f_id)), shape=(int(vehicle_num), int(frame_num))) # i行:车id;j列:时间;元素为i车j时刻的坐标x\n",
    "    sparse_Y = csr_matrix((y, (v_id, f_id)), shape=(int(vehicle_num), int(frame_num))) # i行:车id;j列:时间;元素为i车j时刻的坐标y\n",
    "    sparse_L = csr_matrix((l, (v_id, f_id)), shape=(int(vehicle_num), int(frame_num))) # i行:车id;j列:时间;元素为i车j时刻的下一时刻是否lane change\n",
    "    seq_windows,label = [],[]\n",
    "    for v,graph in enumerate(graph_list):\n",
    "        if graph.data.size==0:\n",
    "            continue\n",
    "        row = np.unique(graph.tocoo().row)\n",
    "        col = np.unique(graph.tocoo().col)\n",
    "        row_start,row_end = row.min(), row.max()+1\n",
    "        col_start,col_end = col.min(), col.max()+1\n",
    "        dense_v = v - row_start\n",
    "        dense_I = graph[row_start:row_end,col_start:col_end].toarray()\n",
    "        dense_x = sparse_X[row_start:row_end,col_start:col_end].toarray()\n",
    "        dense_y = sparse_Y[row_start:row_end,col_start:col_end].toarray()\n",
    "        dense_xy = np.stack((dense_x,dense_y),axis=2) # (vum,total_seq,2)\n",
    "        dense_l = sparse_L[row_start:row_end,col_start:col_end].toarray()\n",
    "        if dense_xy.shape[0]<max_vnum:\n",
    "            padding_num = max_vnum-dense_xy.shape[0]\n",
    "            padding_xy = np.zeros((padding_num,dense_xy.shape[1],dense_xy.shape[2]))\n",
    "            padding_I = np.zeros((padding_num,dense_I.shape[1]))\n",
    "            dense_xy = np.vstack([dense_xy,padding_xy])\n",
    "            dense_I = np.vstack([dense_I,padding_I])\n",
    "            dense_l = np.vstack([dense_l,padding_I])\n",
    "        for i in range(dense_xy.shape[1]): # for loop on sequence dim\n",
    "            if (i+seq_length)*down_sample_rate > dense_xy.shape[1]:\n",
    "                break\n",
    "            window = dense_xy[:,i:i+seq_length*down_sample_rate:down_sample_rate,:] # (vum=30,seq=16,2)\n",
    "            window_l = dense_l[:,i:i+seq_length*down_sample_rate:down_sample_rate] # (vum=30,seq=16)\n",
    "            if sort_func == \"duration\":\n",
    "                dense_seq_I = dense_I[:,i:(i+seq_length)*down_sample_rate:down_sample_rate]\n",
    "                related_score = dense_seq_I.sum(axis=1)\n",
    "                related_score[dense_v] = related_score[dense_v] + 100 # actually 1 is enough\n",
    "                related_rank = np.argsort(-related_score)\n",
    "            elif sort_func == \"distance\":\n",
    "                related_score = ade(window[:,:6,:],window[dense_v,:6,:])\n",
    "                related_rank = np.argsort(related_score)\n",
    "            window = window[related_rank[:max_vnum],:,:]    \n",
    "            seq_windows.append(window)\n",
    "            label.append((window_l[0,6:].sum()>0)*1) # 30,17,2 (0-5),6,(7-16)\n",
    "    if len(seq_windows)==0:\n",
    "        seq_data = None\n",
    "        seq_label = None\n",
    "    else:\n",
    "        seq_data = np.stack(seq_windows)#(n,vum=30,seq=16,2)\n",
    "        seq_label = np.stack(label)\n",
    "    return seq_data,seq_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file_0  processed. time: 10.41. data shape (52619, 30, 17, 2). label shape (52619,).\n",
      "file_1  processed. time: 11.72. data shape (57405, 30, 17, 2). label shape (57405,).\n",
      "file_2  processed. time: 8.57. data shape (45710, 30, 17, 2). label shape (45710,).\n",
      "file_3  processed. time: 12.55. data shape (59114, 30, 17, 2). label shape (59114,).\n",
      "file_4  processed. time: 13.21. data shape (60612, 30, 17, 2). label shape (60612,).\n",
      "file_5  processed. time: 15.90. data shape (68786, 30, 17, 2). label shape (68786,).\n",
      "file_6  processed. time: 7.37. data shape (37417, 30, 17, 2). label shape (37417,).\n",
      "file_7  processed. time: 19.65. data shape (69908, 30, 17, 2). label shape (69908,).\n",
      "file_8  processed. time: 16.01. data shape (64080, 30, 17, 2). label shape (64080,).\n",
      "file_9  processed. time: 7.53. data shape (37657, 30, 17, 2). label shape (37657,).\n",
      "file_10  processed. time: 27.78. data shape (102683, 30, 17, 2). label shape (102683,).\n",
      "file_11  processed. time: 64.25. data shape (181866, 30, 17, 2). label shape (181866,).\n",
      "file_12  processed. time: 66.73. data shape (160709, 30, 17, 2). label shape (160709,).\n",
      "file_13  processed. time: 61.42. data shape (149761, 30, 17, 2). label shape (149761,).\n",
      "file_14  processed. time: 8.92. data shape (47134, 30, 17, 2). label shape (47134,).\n",
      "file_15  processed. time: 10.02. data shape (51232, 30, 17, 2). label shape (51232,).\n",
      "file_16  processed. time: 9.52. data shape (48747, 30, 17, 2). label shape (48747,).\n",
      "file_17  processed. time: 4.75. data shape (28245, 30, 17, 2). label shape (28245,).\n",
      "file_18  processed. time: 9.65. data shape (49295, 30, 17, 2). label shape (49295,).\n",
      "file_19  processed. time: 13.14. data shape (61499, 30, 17, 2). label shape (61499,).\n",
      "file_20  processed. time: 10.87. data shape (53060, 30, 17, 2). label shape (53060,).\n",
      "file_21  processed. time: 5.79. data shape (35981, 30, 17, 2). label shape (35981,).\n",
      "file_22  processed. time: 15.64. data shape (71909, 30, 17, 2). label shape (71909,).\n",
      "file_23  processed. time: 14.91. data shape (66993, 30, 17, 2). label shape (66993,).\n",
      "file_24  processed. time: 107.65. data shape (357294, 30, 17, 2). label shape (357294,).\n",
      "file_25  processed. time: 82.41. data shape (272849, 30, 17, 2). label shape (272849,).\n",
      "file_26  processed. time: 53.13. data shape (147980, 30, 17, 2). label shape (147980,).\n",
      "file_27  processed. time: 42.51. data shape (124823, 30, 17, 2). label shape (124823,).\n",
      "file_28  processed. time: 36.92. data shape (116122, 30, 17, 2). label shape (116122,).\n",
      "file_29  processed. time: 47.83. data shape (135985, 30, 17, 2). label shape (135985,).\n",
      "file_30  processed. time: 40.33. data shape (124833, 30, 17, 2). label shape (124833,).\n",
      "file_31  processed. time: 17.49. data shape (74996, 30, 17, 2). label shape (74996,).\n",
      "file_32  processed. time: 40.18. data shape (123553, 30, 17, 2). label shape (123553,).\n",
      "file_33  processed. time: 26.11. data shape (98166, 30, 17, 2). label shape (98166,).\n",
      "file_34  processed. time: 31.27. data shape (105678, 30, 17, 2). label shape (105678,).\n",
      "file_35  processed. time: 52.15. data shape (148757, 30, 17, 2). label shape (148757,).\n",
      "file_36  processed. time: 31.33. data shape (105946, 30, 17, 2). label shape (105946,).\n",
      "file_37  processed. time: 43.18. data shape (127295, 30, 17, 2). label shape (127295,).\n",
      "file_38  processed. time: 36.30. data shape (109597, 30, 17, 2). label shape (109597,).\n",
      "file_39  processed. time: 42.45. data shape (119901, 30, 17, 2). label shape (119901,).\n",
      "file_40  processed. time: 41.09. data shape (117033, 30, 17, 2). label shape (117033,).\n",
      "file_41  processed. time: 44.26. data shape (124529, 30, 17, 2). label shape (124529,).\n",
      "file_42  processed. time: 39.31. data shape (113926, 30, 17, 2). label shape (113926,).\n",
      "file_43  processed. time: 49.27. data shape (128627, 30, 17, 2). label shape (128627,).\n",
      "file_44  processed. time: 43.99. data shape (123044, 30, 17, 2). label shape (123044,).\n",
      "file_45  processed. time: 65.05. data shape (171964, 30, 17, 2). label shape (171964,).\n",
      "file_46  processed. time: 44.41. data shape (130705, 30, 17, 2). label shape (130705,).\n",
      "file_47  processed. time: 49.28. data shape (132327, 30, 17, 2). label shape (132327,).\n",
      "file_48  processed. time: 33.64. data shape (107342, 30, 17, 2). label shape (107342,).\n",
      "file_49  processed. time: 38.23. data shape (114552, 30, 17, 2). label shape (114552,).\n",
      "file_50  processed. time: 39.46. data shape (116481, 30, 17, 2). label shape (116481,).\n",
      "file_51  processed. time: 34.70. data shape (107715, 30, 17, 2). label shape (107715,).\n",
      "file_52  processed. time: 44.74. data shape (127758, 30, 17, 2). label shape (127758,).\n",
      "file_53  processed. time: 40.70. data shape (116502, 30, 17, 2). label shape (116502,).\n",
      "file_54  processed. time: 37.23. data shape (111609, 30, 17, 2). label shape (111609,).\n",
      "file_55  processed. time: 26.31. data shape (92623, 30, 17, 2). label shape (92623,).\n",
      "file_56  processed. time: 28.72. data shape (96534, 30, 17, 2). label shape (96534,).\n",
      "file_57  processed. time: 6.24. data shape (34985, 30, 17, 2). label shape (34985,).\n",
      "file_58  processed. time: 6.08. data shape (34444, 30, 17, 2). label shape (34444,).\n",
      "file_59  processed. time: 18.70. data shape (76074, 30, 17, 2). label shape (76074,).\n"
     ]
    }
   ],
   "source": [
    "selected_col = [\"frame\",\"id\",\"x\",\"y\",\"laneId\"]\n",
    "total_length = 0\n",
    "for i in range(60):\n",
    "    t1 = time.time()\n",
    "    if i < 9:\n",
    "        path = f\"highD-dataset-v1.0/data/0{i+1}_tracks.csv\"\n",
    "    else:\n",
    "        path = f\"highD-dataset-v1.0/data/{i+1}_tracks.csv\"\n",
    "    data = pd.read_csv(path)\n",
    "    useful_data = data[selected_col]\n",
    "    useful_data = useful_data.sort_values(by=[\"id\",\"frame\"])\n",
    "    label = add_change_label(useful_data)\n",
    "    useful_data[\"label\"] = pd.Series(label)\n",
    "    '''\n",
    "    uni_id = selected_data.id.unique()\n",
    "    mapping_uni_id = np.arange(uni_id.shape[0])\n",
    "    new_uni_id = mapping_uni_id + total_length\n",
    "    total_length += uni_id.shape[0]\n",
    "    id_dict = dict(zip(uni_id, new_uni_id))\n",
    "    new_id = np.vectorize(id_dict.get)(data.id)\n",
    "    selected_data.id = new_id\n",
    "    '''\n",
    "    begin = useful_data[useful_data[\"label\"]==1].index - 20\n",
    "    end = useful_data[useful_data[\"label\"]==1].index + 20\n",
    "    new_label = np.zeros(useful_data.shape[0])\n",
    "    for b,e in zip(begin,end):\n",
    "        new_label[b:e] = 1\n",
    "    useful_data['label'] = new_label\n",
    "    neighbor_graph = getNeighborGraph(useful_data,radius=50)\n",
    "    seq_data,seq_label = graph2seq(useful_data,neighbor_graph,seq_length=17)\n",
    "    t2 = time.time()\n",
    "    print(f\"file_{i}  processed. time: {t2-t1:.2f}. data shape {seq_data.shape}. label shape {seq_label.shape}.\")\n",
    "    file_data = {\"data\":seq_data,\"label\":seq_label}\n",
    "    with open(f\"pickle_data/data_{i}.pkl\",\"wb\") as f:\n",
    "        pkl.dump(file_data,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(234800, 30, 17, 2) (234800,)\n"
     ]
    }
   ],
   "source": [
    "pickle_folder = \"pickle_data\"\n",
    "total_data,total_label = [],[]\n",
    "for file_path in os.listdir(pickle_folder):\n",
    "    if \"_\" not in file_path:\n",
    "        continue\n",
    "    path = os.path.join(pickle_folder,file_path)\n",
    "    with open(path,\"rb\") as f:\n",
    "        pickle_file = pkl.load(f)\n",
    "    label = pickle_file[\"label\"]\n",
    "    data = pickle_file[\"data\"]\n",
    "    pos_data = data[label==1]\n",
    "    neg_data = data[label==0]\n",
    "    pos_data_num = pos_data.shape[0]\n",
    "    total_data.append(pos_data.repeat(10,axis=0))\n",
    "    total_data.append(neg_data[:pos_data_num*10])\n",
    "    total_label.append(np.ones(pos_data_num*10))\n",
    "    total_label.append(np.zeros(pos_data_num*10))\n",
    "data_array = np.vstack(total_data)\n",
    "label_array = np.hstack(total_label)\n",
    "print(data_array.shape, label_array.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22326, 30, 17, 2) (223260, 30, 17, 2)\n"
     ]
    }
   ],
   "source": [
    "pickle_folder = \"pickle_data\"\n",
    "pos_data_list, neg_data_list = [], []\n",
    "for file_path in os.listdir(pickle_folder):\n",
    "    if \"_\" not in file_path:\n",
    "        continue\n",
    "    path = os.path.join(pickle_folder,file_path)\n",
    "    with open(path,\"rb\") as f:\n",
    "        pickle_file = pkl.load(f)\n",
    "    label = pickle_file[\"label\"]\n",
    "    data = pickle_file[\"data\"]\n",
    "    pos_data = data[label==1]\n",
    "    neg_data = data[label==0]\n",
    "    pos_data_num = pos_data.shape[0]\n",
    "    pos_data_list.append(pos_data)\n",
    "    neg_data_list.append(neg_data[:pos_data_num*10])\n",
    "pos_data_array = np.vstack(pos_data_list)\n",
    "neg_data_array = np.vstack(neg_data_list)\n",
    "print(pos_data_array.shape,neg_data_array.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file: data_8.pkl, keep: 950, right: 0, left: 95\n",
      "file: data_3.pkl, keep: 2530, right: 253, left: 11\n",
      "file: data_30.pkl, keep: 550, right: 55, left: 0\n",
      "file: data_17.pkl, keep: 1020, right: 102, left: 0\n",
      "file: data_25.pkl, keep: 260, right: 26, left: 0\n",
      "file: data_44.pkl, keep: 1410, right: 66, left: 141\n",
      "file: data_0.pkl, keep: 1500, right: 0, left: 150\n",
      "file: data_57.pkl, keep: 1340, right: 134, left: 0\n",
      "file: data_38.pkl, keep: 1330, right: 3, left: 133\n",
      "file: data_16.pkl, keep: 1300, right: 107, left: 130\n",
      "file: data_51.pkl, keep: 430, right: 6, left: 43\n",
      "file: data_33.pkl, keep: 700, right: 0, left: 70\n",
      "file: data_49.pkl, keep: 1850, right: 25, left: 185\n",
      "file: data_4.pkl, keep: 1600, right: 87, left: 160\n",
      "file: data_9.pkl, keep: 2810, right: 281, left: 157\n",
      "file: data_46.pkl, keep: 2110, right: 3, left: 211\n",
      "file: data_14.pkl, keep: 540, right: 43, left: 54\n",
      "file: data_13.pkl, keep: 370, right: 26, left: 37\n",
      "file: data_36.pkl, keep: 720, right: 72, left: 6\n",
      "file: data_2.pkl, keep: 940, right: 0, left: 94\n",
      "file: data_28.pkl, keep: 220, right: 1, left: 22\n",
      "file: data_58.pkl, keep: 1160, right: 85, left: 116\n",
      "file: data_18.pkl, keep: 530, right: 0, left: 53\n",
      "file: data_47.pkl, keep: 3200, right: 5, left: 320\n",
      "file: data_39.pkl, keep: 1330, right: 133, left: 76\n",
      "file: data_10.pkl, keep: 190, right: 19, left: 0\n",
      "file: data_53.pkl, keep: 1150, right: 38, left: 115\n",
      "file: data_35.pkl, keep: 4520, right: 452, left: 60\n",
      "file: data_6.pkl, keep: 1530, right: 13, left: 153\n",
      "file: data_52.pkl, keep: 590, right: 59, left: 39\n",
      "file: data_50.pkl, keep: 3370, right: 5, left: 337\n",
      "file: data_59.pkl, keep: 2540, right: 254, left: 197\n",
      "file: data_15.pkl, keep: 460, right: 31, left: 46\n",
      "file: data_42.pkl, keep: 2570, right: 117, left: 257\n",
      "file: data_12.pkl, keep: 50, right: 5, left: 0\n",
      "file: data_23.pkl, keep: 660, right: 66, left: 2\n",
      "file: data_55.pkl, keep: 420, right: 42, left: 8\n",
      "file: data_24.pkl, keep: 2430, right: 243, left: 0\n",
      "file: data_11.pkl, keep: 0, right: 0, left: 0\n",
      "file: data_40.pkl, keep: 1530, right: 0, left: 153\n",
      "file: data_21.pkl, keep: 1390, right: 0, left: 139\n",
      "file: data_31.pkl, keep: 1280, right: 0, left: 128\n",
      "file: data_26.pkl, keep: 170, right: 14, left: 17\n",
      "file: data_41.pkl, keep: 60, right: 6, left: 0\n",
      "file: data_43.pkl, keep: 1770, right: 30, left: 177\n",
      "file: data_56.pkl, keep: 1380, right: 138, left: 39\n",
      "file: data_37.pkl, keep: 9090, right: 61, left: 909\n",
      "file: data_1.pkl, keep: 800, right: 80, left: 36\n",
      "file: data_48.pkl, keep: 5250, right: 93, left: 525\n",
      "file: data_29.pkl, keep: 3730, right: 20, left: 373\n",
      "file: data_32.pkl, keep: 750, right: 47, left: 75\n",
      "file: data_22.pkl, keep: 410, right: 24, left: 41\n",
      "file: data_54.pkl, keep: 2670, right: 0, left: 267\n",
      "file: data_19.pkl, keep: 3820, right: 0, left: 382\n",
      "file: data_20.pkl, keep: 2030, right: 1, left: 203\n",
      "file: data_27.pkl, keep: 2000, right: 23, left: 200\n",
      "file: data_45.pkl, keep: 650, right: 0, left: 65\n",
      "file: data_7.pkl, keep: 3130, right: 313, left: 111\n",
      "file: data_5.pkl, keep: 2540, right: 204, left: 254\n",
      "file: data_34.pkl, keep: 1430, right: 17, left: 143\n",
      "(97060, 30, 17, 2) (3928, 30, 17, 2) (7715, 30, 17, 2)\n"
     ]
    }
   ],
   "source": [
    "pickle_folder = \"pickle_data\"\n",
    "keep_data_list, right_data_list,left_data_list = [], [], []\n",
    "for file_path in os.listdir(pickle_folder):\n",
    "    if \"data_\" not in file_path:\n",
    "        continue\n",
    "    path = os.path.join(pickle_folder,file_path)\n",
    "    with open(path,\"rb\") as f:\n",
    "        pickle_file = pkl.load(f)\n",
    "    label = pickle_file[\"label\"]\n",
    "    data = pickle_file[\"data\"]\n",
    "    left_number = right_number = keep_number = 0\n",
    "    if (label == 1).any():\n",
    "        left_data = data[label==1]\n",
    "        left_data_list.append(left_data)\n",
    "        left_number = left_data.shape[0]\n",
    "    if (label == -1).any():\n",
    "        right_data = data[label==-1]\n",
    "        right_data_list.append(right_data)\n",
    "        right_number = right_data.shape[0]\n",
    "    if (label == 0).any():\n",
    "        keep_number = max(left_number,right_number)*10\n",
    "        keep_data = data[label==0]\n",
    "        keep_data_list.append(keep_data[:keep_number])\n",
    "    print(f\"file: {file_path}, keep: {keep_number}, right: {right_number}, left: {left_number}\")\n",
    "keep_data_array = np.vstack(keep_data_list)\n",
    "right_data_array = np.vstack(right_data_list)\n",
    "left_data_array = np.vstack(left_data_list)\n",
    "print(keep_data_array.shape,right_data_array.shape,left_data_array.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = {\n",
    "    \"right_data\":right_data_array,\n",
    "    \"left_data\":left_data_array,\n",
    "    \"keep_data\":keep_data_array\n",
    "}\n",
    "with open(\"pickle_data/23w10v1_3cls.pkl\",\"wb\") as f:\n",
    "    pkl.dump(_,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = {\n",
    "    \"pos_data\":pos_data_array,\n",
    "    \"neg_data\":neg_data_array,\n",
    "}\n",
    "with open(\"pickle_data/23w10v1.pkl\",\"wb\") as f:\n",
    "    pkl.dump(_,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22326, 30, 17, 2) (223260, 30, 17, 2)\n"
     ]
    }
   ],
   "source": [
    "import pickle as pkl\n",
    "import numpy as np\n",
    "from datatool import train_test_val_split\n",
    "\n",
    "with open(\"pickle_data/23w10v1.pkl\",\"rb\") as f:\n",
    "    data = pkl.load(f)\n",
    "pos_data,neg_data = data['pos_data'],data['neg_data']\n",
    "print(pos_data.shape,neg_data.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(156280, 30, 17, 2) (22320, 30, 17, 2) (44660, 30, 17, 2)\n",
      "(156282, 30, 17, 2) (22326, 30, 17, 2) (44652, 30, 17, 2)\n"
     ]
    }
   ],
   "source": [
    "pos_train, pos_val, pos_test = train_test_val_split(pos_data,test_size=0.2,val_size=0.1,seed=0)\n",
    "neg_train, neg_val, neg_test = train_test_val_split(neg_data,test_size=0.2,val_size=0.1,seed=0)\n",
    "pos_train, pos_val, pos_test = pos_train.repeat(10,axis=0),pos_val.repeat(10,axis=0),pos_test.repeat(10,axis=0)\n",
    "print(pos_train.shape,pos_val.shape,pos_test.shape)\n",
    "print(neg_train.shape,neg_val.shape,neg_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(312562, 30, 17, 2) (312562,)\n",
      "(44646, 30, 17, 2) (44646,)\n",
      "(89312, 30, 17, 2) (89312,)\n"
     ]
    }
   ],
   "source": [
    "X_train, X_val, X_test = np.vstack((pos_train,neg_train)),np.vstack((pos_val,neg_val)),np.vstack((pos_test,neg_test))\n",
    "y_train = np.hstack((np.ones(pos_train.shape[0]),np.zeros(neg_train.shape[0])))\n",
    "y_val = np.hstack((np.ones(pos_val.shape[0]),np.zeros(neg_val.shape[0])))\n",
    "y_test = np.hstack((np.ones(pos_test.shape[0]),np.zeros(neg_test.shape[0])))\n",
    "print(X_train.shape,y_train.shape)\n",
    "print(X_val.shape,y_val.shape)\n",
    "print(X_test.shape,y_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}