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)