101 lines
4.2 KiB
Python
101 lines
4.2 KiB
Python
|
import argparse
|
||
|
|
import datetime
|
||
|
|
import glob
|
||
|
|
import io
|
||
|
|
import ntpath
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
from config import Configuration
|
||
|
|
|
||
|
|
|
||
|
|
def call_training_script(config):
|
||
|
|
cmd = '/home/tomekb/miniconda3/envs/conda3.7/bin/python -u /home/tomekb/yolov3/train.py '
|
||
|
|
cmd += config.train.get_args_string()
|
||
|
|
print("_______ CALLING TRAINING SCRIPT _______")
|
||
|
|
print(cmd)
|
||
|
|
|
||
|
|
os.chdir('..')
|
||
|
|
process = subprocess.Popen(cmd, stdout=subprocess.PIPE,shell=True)
|
||
|
|
for line in io.TextIOWrapper(process.stdout, encoding="utf-8"): # print output of training process to console
|
||
|
|
print(line)
|
||
|
|
return cmd
|
||
|
|
|
||
|
|
def move_training_results_to_experiments_dir(config):
|
||
|
|
training_results_dir_path = os.path.join(config.experiments.dir, datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
|
||
|
|
print("_______ CALLING MOVING RESULTS _______")
|
||
|
|
print(f"MOVING RESUTLS TO {training_results_dir_path}")
|
||
|
|
|
||
|
|
|
||
|
|
os.mkdir(training_results_dir_path)
|
||
|
|
|
||
|
|
weights_path = os.path.join(training_results_dir_path, 'best.pt')
|
||
|
|
shutil.move('/home/tomekb/yolov3/weights/best.pt', weights_path) #move best weights
|
||
|
|
|
||
|
|
names_path = open(config.train.data).readlines()[3].split('=')[-1].rstrip() # read names path from file
|
||
|
|
names_file_name = ntpath.basename(names_path)
|
||
|
|
experiment_names_path = os.path.join(training_results_dir_path, names_file_name)
|
||
|
|
shutil.copy(names_path, experiment_names_path) # copy names to created experiment dir with training results
|
||
|
|
|
||
|
|
tensorboard_dir = './runs'
|
||
|
|
last_modified_tensorboard_dir = max(glob.glob(os.path.join(tensorboard_dir, '*/')), key=os.path.getmtime)
|
||
|
|
shutil.move(last_modified_tensorboard_dir, os.path.join(training_results_dir_path)) #saving related tensorboard dir
|
||
|
|
|
||
|
|
shutil.copy2(config.config_path, training_results_dir_path) #copying configuration yaml
|
||
|
|
|
||
|
|
|
||
|
|
#for test purposes only
|
||
|
|
shutil.copy2('/home/tomekb/yolov3/experiments/1/best.pt', training_results_dir_path)
|
||
|
|
|
||
|
|
return weights_path, experiment_names_path, training_results_dir_path
|
||
|
|
|
||
|
|
|
||
|
|
def call_detection_script(config, weights_path, names_path, dir):
|
||
|
|
detect_output_dir = os.path.join(dir, 'output')
|
||
|
|
cmd = f"""/home/tomekb/miniconda3/envs/conda3.7/bin/python -u /home/tomekb/yolov3/detect.py
|
||
|
|
--cfg {config.train.cfg}
|
||
|
|
--source {config.detect.source}
|
||
|
|
--output {detect_output_dir}
|
||
|
|
--names {names_path}
|
||
|
|
--weights {weights_path}
|
||
|
|
--test-img-size {getattr(config.detect, 'test-img-size')}
|
||
|
|
--conf-thres {getattr(config.detect, 'conf-thres')}
|
||
|
|
--iou-thres {getattr(config.detect, 'iou-thres')}"""
|
||
|
|
cmd += " --save-txt" if getattr(config.detect, 'save-txt') else ""
|
||
|
|
cmd += " --agnostic-nms" if getattr(config.detect, 'agnostic-nms') else ""
|
||
|
|
cmd += " --agument" if getattr(config.detect, 'augment') else ""
|
||
|
|
cmd += f" --device {config.train.device}" if config.train.device else ""
|
||
|
|
|
||
|
|
cmd = " ".join(cmd.split())
|
||
|
|
|
||
|
|
|
||
|
|
print("_______ CALLING DETECTION SCRIPT _______")
|
||
|
|
print(cmd)
|
||
|
|
|
||
|
|
process = subprocess.Popen(cmd, stdout=subprocess.PIPE,shell=True)
|
||
|
|
for line in io.TextIOWrapper(process.stdout, encoding="utf-8"): # print output of process to console
|
||
|
|
print(line)
|
||
|
|
return detect_output_dir
|
||
|
|
|
||
|
|
|
||
|
|
def call_generate_confussion_matrix(detect_output_dir, config, names_path, train_results_dir):
|
||
|
|
labels_dir = getattr(config.confussion_matrix, 'labels-dir')
|
||
|
|
|
||
|
|
cmd = f"node ./our_scripts/generate-confusion-matrix.js {detect_output_dir} {labels_dir} {names_path} > {train_results_dir}/confussion-matrix.tsv"
|
||
|
|
process = subprocess.Popen(cmd, stdout=subprocess.PIPE,shell=True)
|
||
|
|
print("_______ CALLING CONFUSSION MATRIX SCRIPT _______")
|
||
|
|
print(cmd)
|
||
|
|
for line in io.TextIOWrapper(process.stdout, encoding="utf-8"): # print output of process to console
|
||
|
|
print(line)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
config = Configuration()
|
||
|
|
|
||
|
|
train_cmd = call_training_script(config)
|
||
|
|
weights_path, names_path,train_results_dir = move_training_results_to_experiments_dir(config)
|
||
|
|
detect_output_dir = call_detection_script(config, weights_path, names_path,train_results_dir)
|
||
|
|
call_generate_confussion_matrix(detect_output_dir, config, names_path,train_results_dir)
|