def __init__(self, args):
spec_path = os.path.join(args.version, 'spec.yaml')
with open(spec_path) as f:
spec = yaml.load(f)
for key in spec:
setattr(self, key, spec[key])
self.version = args.version
self.ctx = yolo_gluon.get_ctx(args.gpu)
self.export_file = args.version + '/export/'
self.num_downsample = len(self.block_config) + 1
assert args.mode in ['train', 'valid', 'export', 'video']
if args.mode == 'video':
return
self.backup_dir = os.path.join(args.version, 'backup')
self.net = LPDenseNet(self.num_init_features,
self.growth_rate,
self.block_config,
classes=self.LP_num_class)
if args.weight is None:
args.weight = yolo_gluon.get_latest_weight_from(self.backup_dir)
yolo_gluon.init_NN(self.net, args.weight, self.ctx)
if args.mode == 'train':
self.record = args.record
self._init_train()
def __init__(self, args):
spec_path = os.path.join(args.version, 'spec.yaml')
with open(spec_path) as f:
spec = yaml.load(f)
for key in spec:
setattr(self, key, spec[key])
self.ctx = yolo_gluon.get_ctx(args.gpu)
self.export_file = args.version + '/export/'
if args.mode != 'video': # train/valid/export
self.backup_dir = os.path.join(args.version, 'backup')
self.net = LPDenseNet(self.num_init_features,
self.growth_rate,
self.block_config,
classes=self.LP_num_class)
if args.weight is None:
args.weight = yolo_gluon.get_latest_weight_from(
self.backup_dir)
yolo_gluon.init_NN(self.net, args.weight, self.ctx)
if args.mode == 'train':
self.version = args.version
self.record = args.record
self._init_train()
else:
self.tensorrt = args.tensorrt
self.topic = args.topic
self.show = args.show
def _init_net(self, spec, weight):
# (ONNX Error) Do not set num_sync_bn_devices=len(self.ctx)
# because No conversion function for contrib_SyncBatchNorm yet.
self.net = CarNet(spec, num_sync_bn_devices=-1)
if self.use_fp16:
print('cast net to FP16')
self.net.cast('float16')
# self.backup_dir = os.path.join(self.version, 'backup')
self.backup_dir = os.path.join('/media/nolan/SSD1/YOLO_backup/car',
self.version, 'backup')
if weight is None:
weight = yolo_gluon.get_latest_weight_from(self.backup_dir)
yolo_gluon.init_NN(self.net, weight, self.ctx)
def _init_net(self, spec, weight):
# Do not set num_sync_bn_devices=len(self.ctx)
# because No conversion function for contrib_SyncBatchNorm yet.
# (ONNX)
# self.net = CarLPNet(spec, num_sync_bn_devices=len(self.cx))
self.net = CarLPNet(spec, num_sync_bn_devices=-1)
self.num_downsample = len(spec['layers']) - 2 # For LP training
# self.backup_dir = os.path.join(self.version, 'backup')
self.backup_dir = os.path.join(
'/media/nolan/SSD1/YOLO_backup/car_and_LP',
self.version,
'backup')
if weight is None:
weight = yolo_gluon.get_latest_weight_from(self.backup_dir)
yolo_gluon.init_NN(self.net, weight, self.ctx)
def _init_net(self, spec, weight):
from utils import CarDenseNet
num_anchors = len(self.all_anchors[0])
channels_per_anchor = self.slice_point[-1]
channels = num_anchors * channels_per_anchor
self.net = CarDenseNet(self.num_init_features,
self.growth_rate,
self.block_config,
classes=channels - 7,
num_anchors=num_anchors)
if self.use_fp16:
print('cast net to FP16')
self.net.cast('float16')
# self.backup_dir = os.path.join(self.version, 'backup')
self.backup_dir = os.path.join('/media/nolan/SSD1/YOLO_backup/car',
self.version, 'backup')
if weight is None:
weight = yolo_gluon.get_latest_weight_from(self.backup_dir)
yolo_gluon.init_NN(self.net, weight, self.ctx)
rospy.init_node("OCR_node", anonymous=True)
pub = rospy.Publisher('YOLO/OCR', String, queue_size=0)
rospy.Subscriber('/YOLO/clipped_LP', Image, _image_callback)
print('Image Topic: /YOLO/clipped_LP')
print('checkpoint file: %s' % export_file)
r = rospy.Rate(30)
while not rospy.is_shutdown():
r.sleep()
else:
net = OCRDenseNet(num_init_features, growth_rate, block_config, classes=34)
backup_dir = os.path.join(args.version, 'backup')
weight = yolo_gluon.get_latest_weight_from(backup_dir)
yolo_gluon.init_NN(net, weight, ctx)
if args.mode == 'train':
batch_size = 100 * len(ctx)
backward_counter = 0
CE_loss = gluon.loss.SoftmaxCrossEntropyLoss(from_logits=False)
LG_loss = gluon.loss.LogisticLoss(label_format='binary')
L2_loss = gluon.loss.L2Loss()
trainer = gluon.Trainer(net.collect_params(), 'adam',
{'learning_rate': LR})
log_dir = args.version + '/logs'
summary_writer = SummaryWriter(logdir=log_dir, verbose=True)
# prob, score = net(nd.zeros((1,3,160,384), ctx=ctx[0]))