def __init__(self, arg=None, vis=None, id=None):
super().__init__()
self.arg = ARG() if arg is None else arg
self.vis = VIS(save_dir=arg.save.tensorboard) if vis is None else vis
self.id = 0 if id is None else id
self.global_step = 0
self.hyper_parameter = dict()
self.hyper_parameter['lr'] = 10**-1.6
self.dataset = Hand_Dataset(
self.arg,
self.arg.data.dataset.train.root,
im_size=self.arg.model.net.im_size,
relative_path_txt=self.arg.data.dataset.train.txt_path)
self.dataloader = DataLoader(
self.dataset,
batch_size=self.arg.data.dataloader.train.batch_size,
shuffle=self.arg.data.dataloader.train.shuffle,
drop_last=self.arg.data.dataloader.train.drop_last)
# print('train dataloader len = ', len(self.dataloader))
self.net = YoloV3_Micro(class_num=self.arg.model.net.num_class,
anchors_num=self.arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, self.arg.model.net.im_size[0],
self.arg.model.net.im_size[1])
# self.net.forward(input_to_net)
self.vis.model(self.net, input_to_net)
if self.arg.model.weight_path is not None:
self.net.load_state_dict(
one_card_model(torch.load(self.arg.model.weight_path)))
# self.net = torch.nn.DataParallel(self.net)
self.loss_func = YOLOLoss()
# loss_func = torch.nn.DataParallel(loss_func)
self.optimizer = torch.optim.Adam(self.net.parameters(),
lr=self.hyper_parameter['lr'])
# self.lr_decay = ExponentialDecay(0.1, 10000, 1e-4)
self.bbox_format = list()
self.bbox_format.append(
Bbox_Format(self.arg.model.net.im_size,
self.arg.model.net.feature_size[0],
self.arg.model.flt_anchor,
self.arg.model.net.num_class,
self.arg.model.mask_iou_threshold))
self.bbox_format.append(
Bbox_Format(self.arg.model.net.im_size,
self.arg.model.net.feature_size[1],
self.arg.model.flt_anchor,
self.arg.model.net.num_class,
self.arg.model.mask_iou_threshold))
self.valider = Valider(self.arg)
self.tester = Tester(self.arg)
def main():
arg = ARG()
prepare_save_dirs(arg)
vis = VIS(save_dir=arg.save.tensorboard)
# =======================================================
hand_dataset = Hand_Dataset(arg, arg.data.dataset.train.root,
arg.data.dataset.train.relative_txt_path,
arg.model.net.im_size)
train_loader = DataLoader(hand_dataset, batch_size=16, shuffle=False)
vis.iteration_per_epoch = len(train_loader)
net = YoloV3_Micro(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, arg.model.net.im_size[0],
arg.model.net.im_size[1])
net.forward(input_to_net)
vis.model(net, input_to_net)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
loss_func = YOLOLoss(arg.train.device)
loss_func = loss_func.to(arg.train.device)
loss_func = torch.nn.DataParallel(loss_func)
optimizer = torch.optim.Adam(net.parameters(),
lr=arg.optim.lr,
weight_decay=arg.optim.weight_decay)
for epoch in range(1, arg.train.epochs + 1):
vis.epoch_counter = epoch
train(arg, train_loader, net, loss_func, optimizer, epoch, vis)
# if epoch % arg.train.valid_interval == 0:
# valid(arg, valid_loader, model, epoch, vis)
if epoch % arg.train.save_model_epoch_interval == 0:
torch.save(
net.state_dict(),
os.path.join(arg.save.model,
'micro_yolo_epoch{}_weight'.format(epoch)))
def main():
arg = ARG()
prepare_save_dirs(arg)
vis = VIS(save_dir=arg.save.tensorboard)
# =======================================================
train_dataset = Hand_Dataset(arg, arg.data.dataset.train.root, im_size=arg.model.net.im_size, relative_path_txt=arg.data.dataset.train.txt_path)
train_loader = DataLoader(train_dataset, batch_size=arg.data.dataloader.train.batch_size, shuffle=arg.data.dataloader.train.shuffle, drop_last=arg.data.dataloader.train.drop_last)
vis.iteration_per_epoch = len(train_loader)
valid_dataset = Hand_Dataset(arg, arg.data.dataset.valid.root, im_size=arg.model.net.im_size, mode='valid')
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False)
test_dataset = Hand_Dataset(arg, arg.data.dataset.test.root, im_size=arg.model.net.im_size, mode='test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net = YoloV3_Micro(class_num=arg.model.net.num_class, anchors_num=arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, arg.model.net.im_size[0], arg.model.net.im_size[1])
net.forward(input_to_net)
vis.model(net, input_to_net)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
loss_func = YOLOLoss(arg.train.device)
loss_func = loss_func.to(arg.train.device)
loss_func = torch.nn.DataParallel(loss_func)
optimizer = torch.optim.Adam(net.parameters(), lr=arg.optim.lr, weight_decay=arg.optim.weight_decay)
lr_decay = LrSeeker(1e-9, 1)
for epoch in range(1, arg.train.epochs + 1):
vis.epoch_counter = epoch-1
train(arg, train_loader, valid_loader, net, loss_func, optimizer, lr_decay, epoch, vis)
cv2.imshow('', np_im)
cv2.waitKey()
os.makedirs('image_out', exist_ok=True)
cv2.imwrite('image_out/' + str(image_counter) + '.jpg', np_im)
image_counter += 1
rec, prec, ap = det_eval(predict_list, target_dict, 0, 0.5)
return ap
if __name__ == '__main__':
from gear_config.yolo_config.body_micro_yolo_aug1 import ARG
arg = ARG()
test_dataset = Hand_Dataset(arg,
arg.data.dataset.test.root,
im_size=arg.model.net.im_size,
mode='test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net = YoloV3_Micro(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
net = net.to(arg.test.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
tst(arg, test_loader, net, write_image=True)
def tst():
arg = ARG()
net = YoloV3_Tiny(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
net.eval()
bbox_format = list()
bbox_format.append(
Bbox_Format(arg.model.net.im_size, arg.model.net.feature_size[0],
arg.model.anchor, arg.model.net.num_class,
arg.model.mask_iou_threshold))
bbox_format.append(
Bbox_Format(arg.model.net.im_size, arg.model.net.feature_size[1],
arg.model.anchor, arg.model.net.num_class,
arg.model.mask_iou_threshold))
for image_path in os.listdir(arg.test.image_dir):
abs_image_path = os.path.join(arg.test.image_dir, image_path)
image = cv2.imread(abs_image_path)
assert image is not None
torch_im, pad_hwc = yolo_mode_cv2_im_to_torch_im(
image, arg.model.net.im_size)
batch_image = torch_im.unsqueeze(0) # make up as a batch
with torch.no_grad():
net_out = net.forward(batch_image)
predict_bbox_list = list()
confidence_list = list()
for model_out_idx, feature in enumerate(net_out):
batch_predict = model_out_to_model_predict(feature,
num_anchors=len(
arg.model.anchor))
batch_predict_bbox, batch_confidence = bbox_format[
model_out_idx].to_bbox(*batch_predict)
predict_bbox_list.append(
batch_predict_bbox[0]) # notice: batch size == 1
confidence_list.append(
batch_confidence[0]) # notice: batch size == 1
predict_bboxes = torch.cat(predict_bbox_list, 0)
predict_bboxes = cxcywh_to_x1y1x2y2(predict_bboxes)
confidences = torch.cat(confidence_list, 0)
bboxes, confidences = confidence_filter(
predict_bboxes, confidences,
arg.model.out_confidence_filter_threshold)
bboxes, confidences = none_max_suppression(bboxes, confidences,
arg.model.nms_iou_threshold)
np_im = torch_im_to_cv2_im(torch_im).copy()
for box in bboxes:
x1, y1, x2, y2 = list(map(lambda x: int(x), box))
image_pre_show = cv2.rectangle(np_im, (x1, y1), (x2, y2),
(255, 255, 0))
# test_out_dir = arg.test.out_dir
# cv2.imwrite(test_out_dir+'/'+image_path, image_pre_show)
cv2.imshow(image_path, image_pre_show)
cv2.waitKey()
def load_net_weight(self, net):
net.load_state_dict(one_card_model(torch.load(self.arg.model.weight_path)))
def main():
arg = ARG()
prepare_save_dirs(arg)
vis = VIS(save_dir=arg.save.tensorboard)
# =======================================================
train_dataset = Hand_Dataset(
arg,
arg.data.dataset.train.root,
im_size=arg.model.net.im_size,
relative_path_txt=arg.data.dataset.train.txt_path)
train_loader = DataLoader(train_dataset,
batch_size=arg.data.dataloader.train.batch_size,
shuffle=arg.data.dataloader.train.shuffle,
drop_last=arg.data.dataloader.train.drop_last)
vis.iteration_per_epoch = len(train_loader)
valid_dataset = Hand_Dataset(arg,
arg.data.dataset.valid.root,
im_size=arg.model.net.im_size,
mode='valid')
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False)
test_dataset = Hand_Dataset(arg,
arg.data.dataset.test.root,
im_size=arg.model.net.im_size,
mode='test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net = YoloV3_Micro(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, arg.model.net.im_size[0],
arg.model.net.im_size[1])
net.forward(input_to_net)
vis.model(net, input_to_net)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
loss_func = YOLOLoss(arg.train.device)
loss_func = loss_func.to(arg.train.device)
loss_func = torch.nn.DataParallel(loss_func)
optimizer = torch.optim.Adam(net.parameters(), lr=arg.optim.lr)
lr_decay = ExponentialDecay(0.1, 10000, 1e-4)
for epoch in range(1, arg.train.epochs + 1):
vis.epoch_counter = epoch - 1
train(arg, train_loader, valid_loader, net, loss_func, optimizer,
lr_decay, epoch, vis)
if epoch % arg.train.save_model_epoch_interval == 0:
if arg.save.model is not None:
torch.save(
net.state_dict(),
os.path.join(
arg.save.model,
'body_micro_yolo_epoch{}_weight'.format(epoch)))
if epoch % arg.train.test_epoch_interval == 0:
ap = tst(arg, test_loader, net)
vis.line('test/ap', ap)
