yolo_trainer_with_python/train_command_generator.py at master · erdtch/yolo_trainer_with_python · GitHub

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
70
71
72
73
74
75
76
77
78
79
# -*- coding: utf-8 -*-
"""
Created on Wed Jun  5 17:25:58 2019

@author: SarachErudite
"""

import os


"""
Config Training on GPU or Not.
"""

GPU = True  # Set Training False = CPU, True = GPU
GPU_Card = [0,1]  # Number of your GPU [0,1,2,3]

"""
Select yolo version to train your can see at 'darknet/cfg'
"""
is_tiny = True

yolo_cfg = 'custom_yolov3.cfg' # file are contain in 'darknet/cfg'

tiny_yolo_weights = 'yolov3-tiny.weights'

"""
Where to store training-log
"""
train_log = "train.log"

"""

Process code

"""

current_path = os.getcwd()

darknet_dir_path = os.path.join(current_path, 'darknet', '')
darknet_data_path = os.path.join(current_path, 'darknet.data')
pretained_net_path = "./darknet/darknet53.conv.74"
if is_tiny :
    pretained_net_path = "./darknet/darknet53.conv.74"

cfg_path = os.path.join(current_path, 'cfg_data', yolo_cfg)
log_path = os.path.join(current_path, train_log)

command = "./darknet/darknet detector train "
command += darknet_data_path + " "
command += cfg_path + " "
command += pretained_net_path
command += " > " + log_path

if GPU :
    command += " -gpus " + str(GPU_Card).replace('[','').replace(']','').replace(' ','')

if is_tiny :
    tiny_yolo_weights = os.path.join(current_path, 'darknet', tiny_yolo_weights)

    tiny_gen_command = "./darknet/darknet detector partial "
    tiny_gen_command += cfg_path + " "
    tiny_gen_command += tiny_yolo_weights + " "
    tiny_gen_command += "yolov3-tiny.conv.15 15"
    print('\n'*2+'*'*20)
    print('*'*20)
    print("Run this command to generate conv weights for tiny YOLO")
    print('*'*20 + '\n')

    print(tiny_gen_command)

print('\n'*2+'*'*20)
print("Run this command ")
print('*'*20 + '\n')

print(command)

print('\n'+'*'*20 + '\n')
print("Copy and run ...")