class YOLOv3Test(object): def __init__(self, configer): self.configer = configer self.det_visualizer = DetVisualizer(configer) self.det_parser = DetParser(configer) self.det_model_manager = ModelManager(configer) self.test_loader = TestDataLoader(configer) self.device = torch.device('cpu' if self.configer.get('gpu') is None else 'cuda') self.det_net = None self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.det_net.eval() def test(self, test_dir, out_dir): for _, data_dict in enumerate(self.test_loader.get_testloader(test_dir=test_dir)): data_dict['testing'] = True detections = self.det_net(data_dict) meta_list = DCHelper.tolist(data_dict['meta']) batch_detections = self.decode(detections, self.configer, meta_list) for i in range(len(meta_list)): ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'], tool='cv2', mode='BGR') json_dict = self.__get_info_tree(batch_detections[i]) image_canvas = self.det_parser.draw_bboxes(ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) ImageHelper.save(image_canvas, save_path=os.path.join(out_dir, 'vis/{}.png'.format(meta_list[i]['filename']))) Log.info('Json Path: {}'.format(os.path.join(out_dir, 'json/{}.json'.format(meta_list[i]['filename'])))) JsonHelper.save_file(json_dict, save_path=os.path.join(out_dir, 'json/{}.json'.format(meta_list[i]['filename']))) @staticmethod def decode(batch_detections, configer, meta): output = [None for _ in range(len(meta))] for i in range(len(meta)): image_pred = batch_detections[i] image_pred[:, 0] *= meta[i]['ori_img_size'][0] image_pred[:, 1] *= meta[i]['ori_img_size'][1] image_pred[:, 2] *= meta[i]['ori_img_size'][0] image_pred[:, 3] *= meta[i]['ori_img_size'][1] # Filter out confidence scores below threshold image_pred = image_pred[image_pred[:, 4] > configer.get('res', 'val_conf_thre')] # If none are remaining => process next image if image_pred.numel() == 0: continue # Get score and class with highest confidence class_conf, class_pred = torch.max(image_pred[:, 5:5 + configer.get('data', 'num_classes')], 1, keepdim=True) # Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred) detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1) valid_ind = DetHelper.cls_nms(detections[:, :5], labels=class_pred.squeeze(1), max_threshold=configer.get('res', 'nms')['max_threshold'], return_ind=True) output[i] = detections[valid_ind] return output def __get_info_tree(self, detections): json_dict = dict() object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: object_dict = dict() object_dict['bbox'] = [x1.item(), y1.item(), x2.item(), y2.item()] object_dict['label'] = int(cls_pred.item()) object_dict['score'] = float('%.2f' % conf.item()) object_list.append(object_dict) json_dict['objects'] = object_list return json_dict
class YOLOv3Test(object): def __init__(self, configer): self.configer = configer self.blob_helper = BlobHelper(configer) self.det_visualizer = DetVisualizer(configer) self.det_parser = DetParser(configer) self.det_model_manager = ModelManager(configer) self.det_data_loader = DataLoader(configer) self.yolo_target_generator = YOLOTargetGenerator(configer) self.yolo_detection_layer = YOLODetectionLayer(configer) self.device = torch.device( 'cpu' if self.configer.get('gpu') is None else 'cuda') self.det_net = None self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.det_net.eval() def __test_img(self, image_path, json_path, raw_path, vis_path): Log.info('Image Path: {}'.format(image_path)) img = ImageHelper.read_image( image_path, tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) ori_img_bgr = ImageHelper.get_cv2_bgr(img, mode=self.configer.get( 'data', 'input_mode')) inputs = self.blob_helper.make_input(img, input_size=self.configer.get( 'data', 'input_size'), scale=1.0) with torch.no_grad(): inputs = inputs.unsqueeze(0).to(self.device) _, _, detections = self.det_net(inputs) batch_detections = self.decode(detections, self.configer) json_dict = self.__get_info_tree(batch_detections[0], ori_img_bgr) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) ImageHelper.save(ori_img_bgr, raw_path) ImageHelper.save(image_canvas, vis_path) Log.info('Json Path: {}'.format(json_path)) JsonHelper.save_file(json_dict, json_path) return json_dict @staticmethod def decode(batch_pred_bboxes, configer, input_size): box_corner = batch_pred_bboxes.new(batch_pred_bboxes.shape) box_corner[:, :, 0] = batch_pred_bboxes[:, :, 0] - batch_pred_bboxes[:, :, 2] / 2 box_corner[:, :, 1] = batch_pred_bboxes[:, :, 1] - batch_pred_bboxes[:, :, 3] / 2 box_corner[:, :, 2] = batch_pred_bboxes[:, :, 0] + batch_pred_bboxes[:, :, 2] / 2 box_corner[:, :, 3] = batch_pred_bboxes[:, :, 1] + batch_pred_bboxes[:, :, 3] / 2 # clip bounding box box_corner[:, :, 0::2] = box_corner[:, :, 0::2].clamp(min=0, max=1.0) box_corner[:, :, 1::2] = box_corner[:, :, 1::2].clamp(min=0, max=1.0) batch_pred_bboxes[:, :, :4] = box_corner[:, :, :4] batch_pred_bboxes[:, :, 0::2] *= input_size[0] batch_pred_bboxes[:, :, 1::2] *= input_size[1] output = [None for _ in range(len(batch_pred_bboxes))] for image_i, image_pred in enumerate(batch_pred_bboxes): # Filter out confidence scores below threshold conf_mask = (image_pred[:, 4] > configer.get( 'res', 'val_conf_thre')).squeeze() image_pred = image_pred[conf_mask] # If none are remaining => process next image if image_pred.numel() == 0: continue # Get score and class with highest confidence class_conf, class_pred = torch.max( image_pred[:, 5:5 + configer.get('data', 'num_classes')], 1, keepdim=True) # Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred) detections = torch.cat( (image_pred[:, :5], class_conf.float(), class_pred.float()), 1) output[image_i] = DetHelper.cls_nms(detections, labels=class_pred.squeeze(1), max_threshold=configer.get( 'res', 'nms')['max_threshold']) return output def __get_info_tree(self, detections, image_raw, input_size): height, width, _ = image_raw.shape json_dict = dict() object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: object_dict = dict() xmin = x1.cpu().item() / input_size[0] * width ymin = y1.cpu().item() / input_size[1] * height xmax = x2.cpu().item() / input_size[0] * width ymax = y2.cpu().item() / input_size[1] * height object_dict['bbox'] = [xmin, ymin, xmax, ymax] object_dict['label'] = int(cls_pred.cpu().item()) object_dict['score'] = float('%.2f' % conf.cpu().item()) object_list.append(object_dict) json_dict['objects'] = object_list return json_dict def debug(self, vis_dir): count = 0 for i, data_dict in enumerate(self.det_data_loader.get_trainloader()): inputs = data_dict['img'] batch_gt_bboxes = data_dict['bboxes'] batch_gt_labels = data_dict['labels'] input_size = [inputs.size(3), inputs.size(2)] feat_list = list() for stride in self.configer.get('network', 'stride_list'): feat_list.append( torch.zeros((inputs.size(0), 1, input_size[1] // stride, input_size[0] // stride))) targets, _, _ = self.yolo_target_generator(feat_list, batch_gt_bboxes, batch_gt_labels, input_size) targets = targets.to(self.device) anchors_list = self.configer.get('gt', 'anchors_list') output_list = list() be_c = 0 for f_index, anchors in enumerate(anchors_list): feat_stride = self.configer.get('network', 'stride_list')[f_index] fm_size = [ int(round(border / feat_stride)) for border in input_size ] num_c = len(anchors) * fm_size[0] * fm_size[1] output_list.append( targets[:, be_c:be_c + num_c].contiguous().view( targets.size(0), len(anchors), fm_size[1], fm_size[0], -1).permute(0, 1, 4, 2, 3).contiguous().view( targets.size(0), -1, fm_size[1], fm_size[0])) be_c += num_c batch_detections = self.decode( self.yolo_detection_layer(output_list)[2], self.configer, input_size) for j in range(inputs.size(0)): count = count + 1 if count > 20: exit(1) ori_img_bgr = self.blob_helper.tensor2bgr(inputs[j]) json_dict = self.__get_info_tree(batch_detections[j], ori_img_bgr, input_size) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('vis', 'obj_threshold')) cv2.imwrite( os.path.join(vis_dir, '{}_{}_vis.png'.format(i, j)), image_canvas) cv2.imshow('main', image_canvas) cv2.waitKey()
class SingleShotDetectorTest(object): def __init__(self, configer): self.configer = configer self.blob_helper = BlobHelper(configer) self.det_visualizer = DetVisualizer(configer) self.det_parser = DetParser(configer) self.det_model_manager = ModelManager(configer) self.det_data_loader = DataLoader(configer) self.ssd_priorbox_layer = SSDPriorBoxLayer(configer) self.ssd_target_generator = SSDTargetGenerator(configer) self.device = torch.device( 'cpu' if self.configer.get('gpu') is None else 'cuda') self.det_net = None self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.det_net.eval() def __test_img(self, image_path, json_path, raw_path, vis_path): Log.info('Image Path: {}'.format(image_path)) img = ImageHelper.read_image( image_path, tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) ori_img_bgr = ImageHelper.get_cv2_bgr(img, mode=self.configer.get( 'data', 'input_mode')) inputs = self.blob_helper.make_input(img, input_size=self.configer.get( 'test', 'input_size'), scale=1.0) with torch.no_grad(): feat_list, bbox, cls = self.det_net(inputs) batch_detections = self.decode( bbox, cls, self.ssd_priorbox_layer(feat_list, self.configer.get('test', 'input_size')), self.configer, [inputs.size(3), inputs.size(2)]) json_dict = self.__get_info_tree( batch_detections[0], ori_img_bgr, [inputs.size(3), inputs.size(2)]) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) cv2.imwrite(vis_path, image_canvas) cv2.imwrite(raw_path, ori_img_bgr) Log.info('Json Path: {}'.format(json_path)) JsonHelper.save_file(json_dict, json_path) return json_dict @staticmethod def decode(bbox, conf, default_boxes, configer, input_size): loc = bbox if configer.get('phase') != 'debug': conf = F.softmax(conf, dim=-1) default_boxes = default_boxes.unsqueeze(0).repeat(loc.size(0), 1, 1).to(bbox.device) variances = [0.1, 0.2] wh = torch.exp(loc[:, :, 2:] * variances[1]) * default_boxes[:, :, 2:] cxcy = loc[:, :, :2] * variances[ 0] * default_boxes[:, :, 2:] + default_boxes[:, :, :2] boxes = torch.cat([cxcy - wh / 2, cxcy + wh / 2], 2) # [b, 8732,4] batch_size, num_priors, _ = boxes.size() boxes = boxes.unsqueeze(2).repeat(1, 1, configer.get('data', 'num_classes'), 1) boxes = boxes.contiguous().view(boxes.size(0), -1, 4) # clip bounding box boxes[:, :, 0::2] = boxes[:, :, 0::2].clamp(min=0, max=input_size[0] - 1) boxes[:, :, 1::2] = boxes[:, :, 1::2].clamp(min=0, max=input_size[1] - 1) labels = torch.Tensor([ i for i in range(configer.get('data', 'num_classes')) ]).to(boxes.device) labels = labels.view(1, 1, -1, 1).repeat(batch_size, num_priors, 1, 1).contiguous().view(batch_size, -1, 1) max_conf = conf.contiguous().view(batch_size, -1, 1) # max_conf, labels = conf.max(2, keepdim=True) # [b, 8732,1] predictions = torch.cat((boxes, max_conf.float(), labels.float()), 2) output = [None for _ in range(len(predictions))] for image_i, image_pred in enumerate(predictions): ids = labels[image_i].squeeze(1).nonzero().contiguous().view(-1, ) if ids.numel() == 0: continue valid_preds = image_pred[ids] _, order = valid_preds[:, 4].sort(0, descending=True) order = order[:configer.get('nms', 'pre_nms')] valid_preds = valid_preds[order] valid_preds = valid_preds[ valid_preds[:, 4] > configer.get('res', 'val_conf_thre')] if valid_preds.numel() == 0: continue valid_preds = DetHelper.cls_nms( valid_preds[:, :6], labels=valid_preds[:, 5], max_threshold=configer.get('nms', 'max_threshold'), cls_keep_num=configer.get('res', 'cls_keep_num')) _, order = valid_preds[:, 4].sort(0, descending=True) order = order[:configer.get('res', 'max_per_image')] output[image_i] = valid_preds[order] return output def __get_info_tree(self, detections, image_raw, input_size): height, width, _ = image_raw.shape in_width, in_height = input_size json_dict = dict() object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: object_dict = dict() xmin = x1.cpu().item() / in_width * width ymin = y1.cpu().item() / in_height * height xmax = x2.cpu().item() / in_width * width ymax = y2.cpu().item() / in_height * height object_dict['bbox'] = [xmin, ymin, xmax, ymax] object_dict['label'] = int(cls_pred.cpu().item()) - 1 object_dict['score'] = float('%.2f' % conf.cpu().item()) object_list.append(object_dict) json_dict['objects'] = object_list return json_dict def debug(self, vis_dir): count = 0 for i, data_dict in enumerate(self.det_data_loader.get_trainloader()): inputs = data_dict['img'] batch_gt_bboxes = data_dict['bboxes'] batch_gt_labels = data_dict['labels'] input_size = [inputs.size(3), inputs.size(2)] feat_list = list() for stride in self.configer.get('network', 'stride_list'): feat_list.append( torch.zeros((inputs.size(0), 1, input_size[1] // stride, input_size[0] // stride))) bboxes, labels = self.ssd_target_generator(feat_list, batch_gt_bboxes, batch_gt_labels, input_size) eye_matrix = torch.eye(self.configer.get('data', 'num_classes')) labels_target = eye_matrix[labels.view(-1)].view( inputs.size(0), -1, self.configer.get('data', 'num_classes')) batch_detections = self.decode( bboxes, labels_target, self.ssd_priorbox_layer(feat_list, input_size), self.configer, input_size) for j in range(inputs.size(0)): count = count + 1 if count > 20: exit(1) ori_img_bgr = self.blob_helper.tensor2bgr(inputs[j]) self.det_visualizer.vis_default_bboxes( ori_img_bgr, self.ssd_priorbox_layer(feat_list, input_size), labels[j]) json_dict = self.__get_info_tree(batch_detections[j], ori_img_bgr, input_size) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) cv2.imwrite( os.path.join(vis_dir, '{}_{}_vis.png'.format(i, j)), image_canvas) cv2.imshow('main', image_canvas) cv2.waitKey()
class FastRCNNTest(object): def __init__(self, configer): self.configer = configer self.blob_helper = BlobHelper(configer) self.det_visualizer = DetVisualizer(configer) self.det_parser = DetParser(configer) self.det_model_manager = ModelManager(configer) self.test_loader = TestDataLoader(configer) self.roi_sampler = FRROISampler(configer) self.rpn_target_generator = RPNTargetAssigner(configer) self.fr_priorbox_layer = FRPriorBoxLayer(configer) self.fr_roi_generator = FRROIGenerator(configer) self.device = torch.device('cpu' if self.configer.get('gpu') is None else 'cuda') self.det_net = None self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.det_net.eval() def test(self, test_dir, out_dir): for _, data_dict in enumerate(self.test_loader.get_testloader(test_dir=test_dir)): data_dict['testing'] = True out_dict = self.det_net(data_dict) meta_list = DCHelper.tolist(data_dict['meta']) test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = out_dict['test_group'] batch_detections = self.decode(test_roi_locs, test_roi_scores, test_indices_and_rois, test_rois_num, self.configer, meta_list) for i in range(len(meta_list)): ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'], tool='cv2', mode='BGR') json_dict = self.__get_info_tree(batch_detections[i]) image_canvas = self.det_parser.draw_bboxes(ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) ImageHelper.save(image_canvas, save_path=os.path.join(out_dir, 'vis/{}.png'.format(meta_list[i]['filename']))) Log.info('Json Path: {}'.format(os.path.join(out_dir, 'json/{}.json'.format(meta_list[i]['filename'])))) JsonHelper.save_file(json_dict, save_path=os.path.join(out_dir, 'json/{}.json'.format(meta_list[i]['filename']))) @staticmethod def decode(roi_locs, roi_scores, indices_and_rois, test_rois_num, configer, metas): indices_and_rois = indices_and_rois num_classes = configer.get('data', 'num_classes') mean = torch.Tensor(configer.get('roi', 'loc_normalize_mean')).repeat(num_classes)[None] std = torch.Tensor(configer.get('roi', 'loc_normalize_std')).repeat(num_classes)[None] mean = mean.to(roi_locs.device) std = std.to(roi_locs.device) roi_locs = (roi_locs * std + mean) roi_locs = roi_locs.contiguous().view(-1, num_classes, 4) rois = indices_and_rois[:, 1:] rois = rois.contiguous().view(-1, 1, 4).expand_as(roi_locs) wh = torch.exp(roi_locs[:, :, 2:]) * (rois[:, :, 2:] - rois[:, :, :2]) cxcy = roi_locs[:, :, :2] * (rois[:, :, 2:] - rois[:, :, :2]) + (rois[:, :, :2] + rois[:, :, 2:]) / 2 dst_bbox = torch.cat([cxcy - wh / 2, cxcy + wh / 2], 2) # [b, 8732,4] if configer.get('phase') != 'debug': cls_prob = F.softmax(roi_scores, dim=1) else: cls_prob = roi_scores cls_label = torch.LongTensor([i for i in range(num_classes)])\ .contiguous().view(1, num_classes).repeat(indices_and_rois.size(0), 1).to(roi_locs.device) output = [None for _ in range(test_rois_num.size(0))] start_index = 0 for i in range(test_rois_num.size(0)): tmp_dst_bbox = dst_bbox[start_index:start_index+test_rois_num[i]] tmp_dst_bbox[:, :, 0::2] = tmp_dst_bbox[:, :, 0::2].clamp(min=0, max=metas[i]['border_size'][0] - 1) tmp_dst_bbox[:, :, 1::2] = tmp_dst_bbox[:, :, 1::2].clamp(min=0, max=metas[i]['border_size'][1] - 1) tmp_dst_bbox *= (metas[i]['ori_img_size'][0] / metas[i]['border_size'][0]) tmp_cls_prob = cls_prob[start_index:start_index+test_rois_num[i]] tmp_cls_label = cls_label[start_index:start_index+test_rois_num[i]] start_index += test_rois_num[i] mask = (tmp_cls_prob > configer.get('res', 'val_conf_thre')) & (tmp_cls_label > 0) tmp_dst_bbox = tmp_dst_bbox[mask].contiguous().view(-1, 4) if tmp_dst_bbox.numel() == 0: continue tmp_cls_prob = tmp_cls_prob[mask].contiguous().view(-1,).unsqueeze(1) tmp_cls_label = tmp_cls_label[mask].contiguous().view(-1,).unsqueeze(1) valid_preds = torch.cat((tmp_dst_bbox, tmp_cls_prob.float(), tmp_cls_label.float()), 1) valid_ind = DetHelper.cls_nms(valid_preds[:, :5], labels=valid_preds[:, 5], max_threshold=configer.get('res', 'nms')['max_threshold'], return_ind=True) valid_preds = valid_preds[valid_ind] output[i] = valid_preds return output def __get_info_tree(self, detections): json_dict = dict() object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: object_dict = dict() object_dict['bbox'] = [x1.item(), y1.item(), x2.item(), y2.item()] object_dict['label'] = int(cls_pred.cpu().item()) - 1 object_dict['score'] = float('%.2f' % conf.cpu().item()) object_list.append(object_dict) json_dict['objects'] = object_list return json_dict
class FasterRCNN(object): """ The class for Single Shot Detector. Include train, val, test & predict. """ def __init__(self, configer): self.configer = configer self.batch_time = AverageMeter() self.data_time = AverageMeter() self.train_losses = AverageMeter() self.val_losses = AverageMeter() self.det_visualizer = DetVisualizer(configer) self.det_loss_manager = LossManager(configer) self.det_model_manager = ModelManager(configer) self.det_data_loader = DataLoader(configer) self.fr_priorbox_layer = FRPriorBoxLayer(configer) self.det_running_score = DetRunningScore(configer) self.det_net = None self.train_loader = None self.val_loader = None self.optimizer = None self.scheduler = None self.runner_state = dict() self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.optimizer, self.scheduler = Trainer.init(self, self._get_parameters()) self.train_loader = self.det_data_loader.get_trainloader() self.val_loader = self.det_data_loader.get_valloader() def _get_parameters(self): lr_1 = [] lr_2 = [] params_dict = dict(self.det_net.named_parameters()) for key, value in params_dict.items(): if value.requires_grad: if 'bias' in key: lr_2.append(value) else: lr_1.append(value) params = [{ 'params': lr_1, 'lr': self.configer.get('lr', 'base_lr') }, { 'params': lr_2, 'lr': self.configer.get('lr', 'base_lr') * 2., 'weight_decay': 0 }] return params def train(self): """ Train function of every epoch during train phase. """ self.det_net.train() start_time = time.time() # Adjust the learning rate after every epoch. self.runner_state['epoch'] += 1 for i, data_dict in enumerate(self.train_loader): Trainer.update(self) self.data_time.update(time.time() - start_time) # Forward pass. loss = self.det_net(data_dict) loss = loss.mean() self.train_losses.update(loss.item(), len(DCHelper.tolist(data_dict['meta']))) self.optimizer.zero_grad() loss.backward() RunnerHelper.clip_grad(self.det_net, 10.) self.optimizer.step() # Update the vars of the train phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.runner_state['iters'] += 1 # Print the log info & reset the states. if self.runner_state['iters'] % self.configer.get( 'solver', 'display_iter') == 0: Log.info( 'Train Epoch: {0}\tTrain Iteration: {1}\t' 'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t' 'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n' 'Learning rate = {3}\tLoss = {loss.val:.8f} (ave = {loss.avg:.8f})\n' .format(self.runner_state['epoch'], self.runner_state['iters'], self.configer.get('solver', 'display_iter'), RunnerHelper.get_lr(self.optimizer), batch_time=self.batch_time, data_time=self.data_time, loss=self.train_losses)) self.batch_time.reset() self.data_time.reset() self.train_losses.reset() if self.configer.get('lr', 'metric') == 'iters' \ and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'): break # Check to val the current model. if self.runner_state['iters'] % self.configer.get( 'solver', 'test_interval') == 0: self.val() def val(self): """ Validation function during the train phase. """ self.det_net.eval() start_time = time.time() with torch.no_grad(): for j, data_dict in enumerate(self.val_loader): batch_gt_bboxes = DCHelper.tolist(data_dict['bboxes']) batch_gt_labels = DCHelper.tolist(data_dict['labels']) metas = DCHelper.tolist(data_dict['meta']) # Forward pass. loss, test_group = self.det_net(data_dict) # Compute the loss of the train batch & backward. loss = loss.mean() self.val_losses.update(loss.item(), len(metas)) test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = test_group batch_detections = FastRCNNTest.decode(test_roi_locs, test_roi_scores, test_indices_and_rois, test_rois_num, self.configer, metas) batch_pred_bboxes = self.__get_object_list(batch_detections) self.det_running_score.update(batch_pred_bboxes, batch_gt_bboxes, batch_gt_labels) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() RunnerHelper.save_net(self, self.det_net, iters=self.runner_state['iters']) # Print the log info & reset the states. Log.info( 'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t' 'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time, loss=self.val_losses)) Log.info('Val mAP: {}\n'.format(self.det_running_score.get_mAP())) self.det_running_score.reset() self.batch_time.reset() self.val_losses.reset() self.det_net.train() def __get_object_list(self, batch_detections): batch_pred_bboxes = list() for idx, detections in enumerate(batch_detections): object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: xmin = x1.cpu().item() ymin = y1.cpu().item() xmax = x2.cpu().item() ymax = y2.cpu().item() cf = conf.cpu().item() cls_pred = int(cls_pred.cpu().item()) - 1 object_list.append( [xmin, ymin, xmax, ymax, cls_pred, float('%.2f' % cf)]) batch_pred_bboxes.append(object_list) return batch_pred_bboxes
class FastRCNNTest(object): def __init__(self, configer): self.configer = configer self.blob_helper = BlobHelper(configer) self.det_visualizer = DetVisualizer(configer) self.det_parser = DetParser(configer) self.det_model_manager = ModelManager(configer) self.det_data_loader = DataLoader(configer) self.roi_sampler = FRROISampler(configer) self.rpn_target_generator = RPNTargetAssigner(configer) self.fr_priorbox_layer = FRPriorBoxLayer(configer) self.fr_roi_generator = FRROIGenerator(configer) self.device = torch.device( 'cpu' if self.configer.get('gpu') is None else 'cuda') self.det_net = None self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.det_net.eval() def __test_img(self, image_path, json_path, raw_path, vis_path): Log.info('Image Path: {}'.format(image_path)) image = ImageHelper.read_image( image_path, tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) ori_img_bgr = ImageHelper.get_cv2_bgr(image, mode=self.configer.get( 'data', 'input_mode')) width, height = ImageHelper.get_size(image) scale1 = self.configer.get('test', 'resize_bound')[0] / min( width, height) scale2 = self.configer.get('test', 'resize_bound')[1] / max( width, height) scale = min(scale1, scale2) inputs = self.blob_helper.make_input(image, scale=scale) b, c, h, w = inputs.size() border_wh = [w, h] if self.configer.exists('test', 'fit_stride'): stride = self.configer.get('test', 'fit_stride') pad_w = 0 if (w % stride == 0) else stride - (w % stride) # right pad_h = 0 if (h % stride == 0) else stride - (h % stride) # down expand_image = torch.zeros( (b, c, h + pad_h, w + pad_w)).to(inputs.device) expand_image[:, :, 0:h, 0:w] = inputs inputs = expand_image data_dict = dict( img=inputs, meta=DataContainer([[ dict(ori_img_size=ImageHelper.get_size(ori_img_bgr), aug_img_size=border_wh, img_scale=scale, input_size=[inputs.size(3), inputs.size(2)]) ]], cpu_only=True)) with torch.no_grad(): # Forward pass. test_group = self.det_net(data_dict) test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = test_group batch_detections = self.decode(test_roi_locs, test_roi_scores, test_indices_and_rois, test_rois_num, self.configer, DCHelper.tolist(data_dict['meta'])) json_dict = self.__get_info_tree(batch_detections[0], ori_img_bgr, scale=scale) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) cv2.imwrite(vis_path, image_canvas) cv2.imwrite(raw_path, ori_img_bgr) Log.info('Json Path: {}'.format(json_path)) JsonHelper.save_file(json_dict, json_path) return json_dict @staticmethod def decode(roi_locs, roi_scores, indices_and_rois, test_rois_num, configer, metas): indices_and_rois = indices_and_rois num_classes = configer.get('data', 'num_classes') mean = torch.Tensor(configer.get( 'roi', 'loc_normalize_mean')).repeat(num_classes)[None] std = torch.Tensor(configer.get( 'roi', 'loc_normalize_std')).repeat(num_classes)[None] mean = mean.to(roi_locs.device) std = std.to(roi_locs.device) roi_locs = (roi_locs * std + mean) roi_locs = roi_locs.contiguous().view(-1, num_classes, 4) # roi_locs = roi_locs[:,:, [1, 0, 3, 2]] rois = indices_and_rois[:, 1:] rois = rois.contiguous().view(-1, 1, 4).expand_as(roi_locs) wh = torch.exp(roi_locs[:, :, 2:]) * (rois[:, :, 2:] - rois[:, :, :2]) cxcy = roi_locs[:, :, :2] * (rois[:, :, 2:] - rois[:, :, :2]) + ( rois[:, :, :2] + rois[:, :, 2:]) / 2 dst_bbox = torch.cat([cxcy - wh / 2, cxcy + wh / 2], 2) # [b, 8732,4] if configer.get('phase') != 'debug': cls_prob = F.softmax(roi_scores, dim=1) else: cls_prob = roi_scores cls_label = torch.LongTensor([i for i in range(num_classes)])\ .contiguous().view(1, num_classes).repeat(indices_and_rois.size(0), 1).to(roi_locs.device) output = [None for _ in range(test_rois_num.size(0))] start_index = 0 for i in range(test_rois_num.size(0)): # batch_index = (indices_and_rois[:, 0] == i).nonzero().contiguous().view(-1,) # tmp_dst_bbox = dst_bbox[batch_index] # tmp_cls_prob = cls_prob[batch_index] # tmp_cls_label = cls_label[batch_index] tmp_dst_bbox = dst_bbox[start_index:start_index + test_rois_num[i]] # clip bounding box tmp_dst_bbox[:, :, 0::2] = tmp_dst_bbox[:, :, 0::2].clamp( min=0, max=metas[i]['border_size'][0] - 1) tmp_dst_bbox[:, :, 1::2] = tmp_dst_bbox[:, :, 1::2].clamp( min=0, max=metas[i]['border_size'][1] - 1) tmp_cls_prob = cls_prob[start_index:start_index + test_rois_num[i]] tmp_cls_label = cls_label[start_index:start_index + test_rois_num[i]] start_index += test_rois_num[i] mask = (tmp_cls_prob > configer.get( 'res', 'val_conf_thre')) & (tmp_cls_label > 0) tmp_dst_bbox = tmp_dst_bbox[mask].contiguous().view(-1, 4) if tmp_dst_bbox.numel() == 0: continue tmp_cls_prob = tmp_cls_prob[mask].contiguous().view( -1, ).unsqueeze(1) tmp_cls_label = tmp_cls_label[mask].contiguous().view( -1, ).unsqueeze(1) valid_preds = torch.cat( (tmp_dst_bbox, tmp_cls_prob.float(), tmp_cls_label.float()), 1) output[i] = DetHelper.cls_nms(valid_preds, labels=valid_preds[:, 5], max_threshold=configer.get( 'nms', 'max_threshold')) return output def __get_info_tree(self, detections, image_raw, scale=1.0): height, width, _ = image_raw.shape json_dict = dict() object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: object_dict = dict() xmin = min(x1.cpu().item() / scale, width - 1) ymin = min(y1.cpu().item() / scale, height - 1) xmax = min(x2.cpu().item() / scale, width - 1) ymax = min(y2.cpu().item() / scale, height - 1) object_dict['bbox'] = [xmin, ymin, xmax, ymax] object_dict['label'] = int(cls_pred.cpu().item()) - 1 object_dict['score'] = float('%.2f' % conf.cpu().item()) object_list.append(object_dict) json_dict['objects'] = object_list return json_dict def debug(self, vis_dir): count = 0 for i, data_dict in enumerate(self.det_data_loader.get_trainloader()): feat_list = list() input_size = data_dict['meta'][0]['input_size'] for stride in self.configer.get('rpn', 'stride_list'): feat_list.append( torch.zeros( (data_dict['img'].size(0), 1, input_size[1] // stride, input_size[0] // stride))) gt_rpn_locs, gt_rpn_labels = self.rpn_target_generator( feat_list, data_dict['bboxes'], data_dict['meta']) eye_matrix = torch.eye(2) gt_rpn_labels[gt_rpn_labels == -1] = 0 gt_rpn_scores = eye_matrix[gt_rpn_labels.view(-1)].view( data_dict['img'].size(0), -1, 2) test_indices_and_rois, _ = self.fr_roi_generator( feat_list, gt_rpn_locs, gt_rpn_scores, self.configer.get('rpn', 'n_test_pre_nms'), self.configer.get('rpn', 'n_test_post_nms'), data_dict['meta']) sample_rois, gt_roi_locs, gt_roi_labels = self.roi_sampler( test_indices_and_rois, data_dict['bboxes'], data_dict['labels'], data_dict['meta']) self.det_visualizer.vis_rois(data_dict['img'], sample_rois[gt_roi_labels > 0]) gt_cls_roi_locs = torch.zeros( (gt_roi_locs.size(0), self.configer.get('data', 'num_classes'), 4)) gt_cls_roi_locs[torch.arange(0, sample_rois.size(0)).long(), gt_roi_labels.long()] = gt_roi_locs gt_cls_roi_locs = gt_cls_roi_locs.contiguous().view( -1, 4 * self.configer.get('data', 'num_classes')) eye_matrix = torch.eye(self.configer.get('data', 'num_classes')) gt_roi_scores = eye_matrix[gt_roi_labels.view(-1)].view( gt_roi_labels.size(0), self.configer.get('data', 'num_classes')) test_rois_num = torch.zeros((len(data_dict['bboxes']), )).long() for batch_id in range(len(data_dict['bboxes'])): batch_index = ( sample_rois[:, 0] == batch_id).nonzero().contiguous().view( -1, ) test_rois_num[batch_id] = batch_index.numel() batch_detections = FastRCNNTest.decode(gt_cls_roi_locs, gt_roi_scores, sample_rois, test_rois_num, self.configer, data_dict['meta']) for j in range(data_dict['img'].size(0)): count = count + 1 if count > 20: exit(1) ori_img_bgr = self.blob_helper.tensor2bgr(data_dict['img'][j]) self.det_visualizer.vis_default_bboxes( ori_img_bgr, self.fr_priorbox_layer(feat_list, input_size), gt_rpn_labels[j]) json_dict = self.__get_info_tree(batch_detections[j], ori_img_bgr) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) cv2.imwrite( os.path.join(vis_dir, '{}_{}_vis.png'.format(i, j)), image_canvas) cv2.imshow('main', image_canvas) cv2.waitKey()
class SingleShotDetector(object): """ The class for Single Shot Detector. Include train, val, test & predict. """ def __init__(self, configer): self.configer = configer self.batch_time = AverageMeter() self.data_time = AverageMeter() self.train_losses = AverageMeter() self.val_losses = AverageMeter() self.det_visualizer = DetVisualizer(configer) self.det_loss_manager = LossManager(configer) self.det_model_manager = ModelManager(configer) self.det_data_loader = DataLoader(configer) self.ssd_target_generator = SSDTargetGenerator(configer) self.ssd_priorbox_layer = SSDPriorBoxLayer(configer) self.det_running_score = DetRunningScore(configer) self.det_net = None self.train_loader = None self.val_loader = None self.optimizer = None self.scheduler = None self.runner_state = dict() self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.optimizer, self.scheduler = Trainer.init(self, self._get_parameters()) self.train_loader = self.det_data_loader.get_trainloader() self.val_loader = self.det_data_loader.get_valloader() self.det_loss = self.det_loss_manager.get_det_loss() def _get_parameters(self): lr_1 = [] lr_10 = [] params_dict = dict(self.det_net.named_parameters()) for key, value in params_dict.items(): if 'backbone' not in key: lr_10.append(value) else: lr_1.append(value) params = [{ 'params': lr_1, 'lr': self.configer.get('lr', 'base_lr') }, { 'params': lr_10, 'lr': self.configer.get('lr', 'base_lr') }] return params def train(self): """ Train function of every epoch during train phase. """ self.det_net.train() start_time = time.time() # Adjust the learning rate after every epoch. self.runner_state['epoch'] += 1 # data_tuple: (inputs, heatmap, maskmap, vecmap) for i, data_dict in enumerate(self.train_loader): Trainer.update(self, backbone_list=(0, )) inputs = data_dict['img'] batch_gt_bboxes = data_dict['bboxes'] batch_gt_labels = data_dict['labels'] # Change the data type. inputs = RunnerHelper.to_device(self, inputs) self.data_time.update(time.time() - start_time) # Forward pass. outputs = self.det_net(inputs) if self.configer.get('network', 'gathered'): feat_list = outputs[0] else: feat_list = outputs[0][0] bboxes, labels = self.ssd_target_generator( feat_list, batch_gt_bboxes, batch_gt_labels, [inputs.size(3), inputs.size(2)]) bboxes, labels = RunnerHelper.to_device(bboxes, labels) # Compute the loss of the train batch & backward. loss = self.det_loss(outputs, bboxes, labels, gathered=self.configer.get( 'network', 'gathered')) self.train_losses.update(loss.item(), inputs.size(0)) self.optimizer.zero_grad() loss.backward() self.optimizer.step() # Update the vars of the train phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.configer.plus_one('iters') # Print the log info & reset the states. if self.configer.get('iters') % self.configer.get( 'solver', 'display_iter') == 0: Log.info( 'Train Epoch: {0}\tTrain Iteration: {1}\t' 'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t' 'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n' 'Learning rate = {3}\tLoss = {loss.val:.8f} (ave = {loss.avg:.8f})\n' .format(self.runner_state['epoch'], self.runner_state['iters'], self.configer.get('solver', 'display_iter'), RunnerHelper.get_lr(self.optimizer), batch_time=self.batch_time, data_time=self.data_time, loss=self.train_losses)) self.batch_time.reset() self.data_time.reset() self.train_losses.reset() if self.configer.get('lr', 'metric') == 'iters' \ and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'): break # Check to val the current model. if self.runner_state['iters'] % self.configer.get( 'solver', 'test_interval') == 0: self.val() def val(self): """ Validation function during the train phase. """ self.det_net.eval() start_time = time.time() with torch.no_grad(): for j, data_dict in enumerate(self.val_loader): inputs = data_dict['img'] batch_gt_bboxes = data_dict['bboxes'] batch_gt_labels = data_dict['labels'] inputs = RunnerHelper.to_device(self, inputs) input_size = [inputs.size(3), inputs.size(2)] # Forward pass. outputs = self.det_net(inputs) feat_list, loc, cls = RunnerHelper.gather(self, outputs) bboxes, labels = self.ssd_target_generator( feat_list, batch_gt_bboxes, batch_gt_labels, input_size) bboxes, labels = RunnerHelper.to_device(bboxes, labels) # Compute the loss of the val batch. loss = self.det_loss(outputs, bboxes, labels, gathered=self.configer.get( 'network', 'gathered')) self.val_losses.update(loss.item(), inputs.size(0)) batch_detections = SingleShotDetectorTest.decode( loc, cls, self.ssd_priorbox_layer(feat_list, input_size), self.configer, input_size) batch_pred_bboxes = self.__get_object_list(batch_detections) # batch_pred_bboxes = self._get_gt_object_list(batch_gt_bboxes, batch_gt_labels) self.det_running_score.update(batch_pred_bboxes, batch_gt_bboxes, batch_gt_labels) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() RunnerHelper.save_net(self, self.det_net, iters=self.runner_state['iters']) # Print the log info & reset the states. Log.info( 'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t' 'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time, loss=self.val_losses)) Log.info('Val mAP: {}'.format(self.det_running_score.get_mAP())) self.det_running_score.reset() self.batch_time.reset() self.val_losses.reset() self.det_net.train() def _get_gt_object_list(self, batch_gt_bboxes, batch_gt_labels): batch_pred_bboxes = list() for i in range(len(batch_gt_bboxes)): object_list = list() if batch_gt_bboxes[i].numel() > 0: for j in range(batch_gt_bboxes[i].size(0)): object_list.append([ batch_gt_bboxes[i][j][0].item(), batch_gt_bboxes[i][j][1].item(), batch_gt_bboxes[i][j][2].item(), batch_gt_bboxes[i][j][3].item(), batch_gt_labels[i][j].item(), 1.0 ]) batch_pred_bboxes.append(object_list) return batch_pred_bboxes def __get_object_list(self, batch_detections): batch_pred_bboxes = list() for idx, detections in enumerate(batch_detections): object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: xmin = x1.cpu().item() ymin = y1.cpu().item() xmax = x2.cpu().item() ymax = y2.cpu().item() cf = conf.cpu().item() cls_pred = cls_pred.cpu().item() - 1 object_list.append([ xmin, ymin, xmax, ymax, int(cls_pred), float('%.2f' % cf) ]) batch_pred_bboxes.append(object_list) return batch_pred_bboxes
class SingleShotDetectorTest(object): def __init__(self, configer): self.configer = configer self.blob_helper = BlobHelper(configer) self.det_visualizer = DetVisualizer(configer) self.det_parser = DetParser(configer) self.det_model_manager = ModelManager(configer) self.test_loader = TestDataLoader(configer) self.device = torch.device( 'cpu' if self.configer.get('gpu') is None else 'cuda') self.det_net = None self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.det_net.eval() def test(self, test_dir, out_dir): for _, data_dict in enumerate( self.test_loader.get_testloader(test_dir=test_dir)): data_dict['testing'] = True loc, conf = self.det_net(data_dict) meta_list = DCHelper.tolist(data_dict['meta']) batch_detections = self.decode(loc, conf, self.configer, meta_list) for i in range(len(meta_list)): ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'], tool='cv2', mode='BGR') json_dict = self.__get_info_tree(batch_detections[i]) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) ImageHelper.save(image_canvas, save_path=os.path.join( out_dir, 'vis/{}.png'.format( meta_list[i]['filename']))) Log.info('Json Path: {}'.format( os.path.join( out_dir, 'json/{}.json'.format(meta_list[i]['filename'])))) JsonHelper.save_file(json_dict, save_path=os.path.join( out_dir, 'json/{}.json'.format( meta_list[i]['filename']))) @staticmethod def decode(loc, conf, configer, meta): batch_size, num_priors, _ = loc.size() loc = loc.unsqueeze(2).repeat(1, 1, configer.get('data', 'num_classes'), 1) loc = loc.contiguous().view(loc.size(0), -1, 4) labels = torch.Tensor([ i for i in range(configer.get('data', 'num_classes')) ]).to(loc.device) labels = labels.view(1, 1, -1, 1).repeat(batch_size, num_priors, 1, 1).contiguous().view(batch_size, -1, 1) conf = conf.contiguous().view(batch_size, -1, 1) # max_conf, labels = conf.max(2, keepdim=True) # [b, 8732,1] predictions = torch.cat((loc.float(), conf.float(), labels.float()), 2) output = [None for _ in range(len(predictions))] for i, image_pred in enumerate(predictions): image_pred[:, 0] *= meta[i]['ori_img_size'][0] image_pred[:, 1] *= meta[i]['ori_img_size'][1] image_pred[:, 2] *= meta[i]['ori_img_size'][0] image_pred[:, 3] *= meta[i]['ori_img_size'][1] ids = labels[i].squeeze(1).nonzero().contiguous().view(-1, ) if ids.numel() == 0: continue valid_preds = image_pred[ids] _, order = valid_preds[:, 4].sort(0, descending=True) order = order[:configer.get('res', 'nms')['pre_nms']] valid_preds = valid_preds[order] valid_preds = valid_preds[ valid_preds[:, 4] > configer.get('res', 'val_conf_thre')] if valid_preds.numel() == 0: continue valid_ind = DetHelper.cls_nms( valid_preds[:, :5], labels=valid_preds[:, 5], max_threshold=configer.get('res', 'nms')['max_threshold'], cls_keep_num=configer.get('res', 'cls_keep_num'), return_ind=True) valid_preds = valid_preds[valid_ind] _, order = valid_preds[:, 4].sort(0, descending=True) order = order[:configer.get('res', 'max_per_image')] output[i] = valid_preds[order] return output def __get_info_tree(self, detections): json_dict = dict() object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: object_dict = dict() object_dict['bbox'] = [ x1.item(), y1.item(), x2.item(), y2.item() ] object_dict['label'] = int(cls_pred.cpu().item()) - 1 object_dict['score'] = float('%.2f' % conf.cpu().item()) object_list.append(object_dict) json_dict['objects'] = object_list return json_dict
class SingleShotDetector(object): """ The class for Single Shot Detector. Include train, val, test & predict. """ def __init__(self, configer): self.configer = configer self.batch_time = AverageMeter() self.data_time = AverageMeter() self.train_losses = AverageMeter() self.val_losses = AverageMeter() self.det_visualizer = DetVisualizer(configer) self.det_model_manager = ModelManager(configer) self.det_data_loader = DataLoader(configer) self.det_running_score = DetRunningScore(configer) self.det_net = None self.train_loader = None self.val_loader = None self.optimizer = None self.scheduler = None self.runner_state = dict() self._init_model() def _init_model(self): self.det_net = self.det_model_manager.object_detector() self.det_net = RunnerHelper.load_net(self, self.det_net) self.optimizer, self.scheduler = Trainer.init( self._get_parameters(), self.configer.get('solver')) self.train_loader = self.det_data_loader.get_trainloader() self.val_loader = self.det_data_loader.get_valloader() self.det_loss = self.det_model_manager.get_det_loss() def _get_parameters(self): lr_1 = [] lr_10 = [] params_dict = dict(self.det_net.named_parameters()) for key, value in params_dict.items(): if 'backbone' not in key: lr_10.append(value) else: lr_1.append(value) params = [{ 'params': lr_1, 'lr': self.configer.get('solver', 'lr')['base_lr'] }, { 'params': lr_10, 'lr': self.configer.get('solver', 'lr')['base_lr'] * 4.0 }] return params def train(self): """ Train function of every epoch during train phase. """ self.det_net.train() start_time = time.time() # Adjust the learning rate after every epoch. self.runner_state['epoch'] += 1 # data_tuple: (inputs, heatmap, maskmap, vecmap) for i, data_dict in enumerate(self.train_loader): Trainer.update(self, backbone_list=(0, ), backbone_lr_list=(self.configer.get( 'solver', 'lr')['base_lr'], ), solver_dict=self.configer.get('solver')) self.data_time.update(time.time() - start_time) # Forward pass. out_dict = self.det_net(data_dict) loss = out_dict['loss'].mean() self.train_losses.update(loss.item(), len(DCHelper.tolist(data_dict['meta']))) self.optimizer.zero_grad() loss.backward() RunnerHelper.clip_grad(self.det_net, 10.) self.optimizer.step() # Update the vars of the train phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.runner_state['iters'] += 1 # Print the log info & reset the states. if self.runner_state['iters'] % self.configer.get( 'solver', 'display_iter') == 0: Log.info( 'Train Epoch: {0}\tTrain Iteration: {1}\t' 'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t' 'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n' 'Learning rate = {3}\tLoss = {loss.val:.8f} (ave = {loss.avg:.8f})\n' .format(self.runner_state['epoch'], self.runner_state['iters'], self.configer.get('solver', 'display_iter'), RunnerHelper.get_lr(self.optimizer), batch_time=self.batch_time, data_time=self.data_time, loss=self.train_losses)) self.batch_time.reset() self.data_time.reset() self.train_losses.reset() if self.configer.get('solver', 'lr')['metric'] == 'iters' \ and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'): break # Check to val the current model. if self.runner_state['iters'] % self.configer.get( 'solver', 'test_interval') == 0: self.val() def val(self): """ Validation function during the train phase. """ self.det_net.eval() start_time = time.time() with torch.no_grad(): for j, data_dict in enumerate(self.val_loader): # Forward pass. out_dict = self.det_net(data_dict) loss = out_dict['loss'].mean() # Compute the loss of the val batch. self.val_losses.update(loss.item(), len(DCHelper.tolist(data_dict['meta']))) batch_detections = SingleShotDetectorTest.decode( out_dict['loc'], out_dict['conf'], self.configer, DCHelper.tolist(data_dict['meta'])) batch_pred_bboxes = self.__get_object_list(batch_detections) # batch_pred_bboxes = self._get_gt_object_list(batch_gt_bboxes, batch_gt_labels) self.det_running_score.update(batch_pred_bboxes, [ item['ori_bboxes'] for item in DCHelper.tolist(data_dict['meta']) ], [ item['ori_labels'] for item in DCHelper.tolist(data_dict['meta']) ]) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() RunnerHelper.save_net(self, self.det_net, iters=self.runner_state['iters']) # Print the log info & reset the states. Log.info( 'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t' 'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time, loss=self.val_losses)) Log.info('Val mAP: {}'.format(self.det_running_score.get_mAP())) self.det_running_score.reset() self.batch_time.reset() self.val_losses.reset() self.det_net.train() def __get_object_list(self, batch_detections): batch_pred_bboxes = list() for idx, detections in enumerate(batch_detections): object_list = list() if detections is not None: for x1, y1, x2, y2, conf, cls_pred in detections: cf = float('%.2f' % conf.item()) cls_pred = int(cls_pred.cpu().item() - 1) object_list.append([ x1.item(), y1.item(), x2.item(), y2.item(), cls_pred, cf ]) batch_pred_bboxes.append(object_list) return batch_pred_bboxes