-
Notifications
You must be signed in to change notification settings - Fork 1
/
train_helen.py
43 lines (33 loc) · 1.6 KB
/
train_helen.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
from lns.common.preprocess import Preprocessor
import os
dataset_scale = Preprocessor.preprocess('ScaleLights')
dataset_utias = Preprocessor.preprocess('ScaleLights_New_Utias')
dataset_youtube = Preprocessor.preprocess('ScaleLights_New_Youtube')
dataset_scale_utias = dataset_scale.__add__(dataset_utias)
dataset_all = dataset_scale_utias.__add__(dataset_youtube)
dataset_all = dataset_all.merge_classes({
"green": ["goLeft", "Green", "GreenLeft", "GreenStraightRight", "go", "GreenStraightLeft", "GreenRight", "GreenStraight", "3-green", "4-green", "5-green"],
"yellow": ["warning", "Yellow", "warningLeft", "3-yellow", "4-yellow", "5-yellow"],
"red": ["stop", "stopLeft", "RedStraightLeft", "Red", "RedLeft", "RedStraight", "RedRight", "3-red", "4-red", "5-red"],
"off": ["OFF", "off", "3-off", "3-other", "4-off", "4-other", "5-off", "5-other"]
})
# from lns.yolo.train import YoloTrainer
# trainer = YoloTrainer('darknet25_640_416_helen',dataset_all)
# os.environ["CUDA_VISIBLE_DEVICES"]="2"
from lns.squeezedet.train import SqueezedetTrainer
trainer = SqueezedetTrainer('brian_lights_trial2',dataset_all,load=False)
trainer.train()
#### train haar on lisa and scale
#from lns.common.preprocess import Preprocessor
#from lns.haar.train import HaarTrainer
#dataset_lisa = Preprocessor.preprocess('lisa_signs')
#dataset_scale = Preprocessor.preprocess('ScaleSigns', force=True)
# print(dataset_scale)
#dataset = dataset_lisa.__add__(dataset_scale)
# dataset = dataset.merge_classes({
# })
# print(dataset.classes)
# exit()
# trainer = HaarTrainer('helen_test_all',dataset)
# trainer.setup()
# trainer.train()