-
Notifications
You must be signed in to change notification settings - Fork 10
/
preprocess.py
69 lines (58 loc) · 2.65 KB
/
preprocess.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import tqdm
import numpy as np
import os
import sys
import json
def eprint(*args, **kwargs):
print(*args, file=sys.stderr, **kwargs)
def load_origin_data(data_dir, app_num):
datas = [[] for _ in range(app_num)]
filenames = [filename for filename in os.listdir(data_dir) \
if os.path.isfile(os.path.join(data_dir, filename)) and filename.split(".")[1] == "num"]
filenames = sorted(filenames)
for app in tqdm.tqdm(range(app_num), ascii=True, desc='[Load Data]'):
with open(os.path.join(data_dir, filenames[app])) as fp:
for line in fp:
length = line.strip().split(';')[1].strip().split('\t')
length = [int(ix) for ix in length]
status = line.strip().split(';')[0].strip().split('\t')
status = [int(ix) for ix in status]
datas[app].append({'label': app, 'flow': length, 'status': status, 'lo': length.copy()})
return datas
def _transform(datas, limit, max_packet):
data_trans = [[] for _ in range(len(datas))]
for app in tqdm.tqdm(range(len(datas)), ascii=True, desc='[Transform]'):
app_data = datas[app]
for idx, example in enumerate(app_data):
flow = example['flow']
if len(flow) < limit:
print(flow)
continue
flow = [ix if ix <= max_packet else max_packet for ix in flow]
data_trans[app].append(
{'label': example['label'], 'flow': flow, 'status': example['status'], 'lo': example['lo'],
'id': str(app) + '-' + str(idx)}
)
return data_trans
def split_train_and_dev(datas, ratio=0.8, keep_ratio=1):
train, dev = [], []
for app_data in tqdm.tqdm(datas, ascii=True, desc='[Split]'):
is_keep = np.random.rand(len(app_data)) <= keep_ratio
is_train = np.random.rand(len(app_data)) <= ratio
for example, kp, tr in zip(app_data, is_keep, is_train):
if kp and tr:
train.append(example)
elif kp and not tr:
dev.append(example)
np.random.shuffle(train)
np.random.shuffle(dev)
return train, dev
def preprocess(config):
eprint('Generate train and test.')
origin = load_origin_data(config.data_dir, config.class_num)
length = _transform(origin, config.min_length, config.max_packet_length)
train, test = split_train_and_dev(length, config.split_ratio, config.keep_ratio)
with open(config.train_json, 'w') as fp:
json.dump(train, fp, indent=1)
with open(config.test_json, 'w') as fp:
json.dump(test, fp, indent=1)