class ADE20KEvaluator(object): def __init__(self, configer): self.configer = configer self.seg_running_score = RunningScore(configer) def relabel(self, labelmap): return (labelmap - 1).astype(np.uint8) def evaluate(self, pred_dir, gt_dir): img_cnt = 0 for filename in os.listdir(pred_dir): pred_path = os.path.join(pred_dir, filename) gt_path = os.path.join(gt_dir, filename) predmap = ImageHelper.img2np( ImageHelper.read_image(pred_path, tool='pil', mode='P')) gtmap = ImageHelper.img2np( ImageHelper.read_image(gt_path, tool='pil', mode='P')) predmap = self.relabel(predmap) gtmap = self.relabel(gtmap) self.seg_running_score.update(predmap[np.newaxis, :, :], gtmap[np.newaxis, :, :]) img_cnt += 1 Log.info('Evaluate {} images'.format(img_cnt)) Log.info('mIOU: {}'.format(self.seg_running_score.get_mean_iou())) Log.info('Pixel ACC: {}'.format( self.seg_running_score.get_pixel_acc()))
class COCOStuffEvaluator(object): def __init__(self, configer): self.configer = configer self.seg_running_score = RunningScore(configer) self.id_to_trainid = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 11, 13: 12, 14: 13, 15: 14, 16: 15, 17: 16, 18: 17, 19: 18, 20: 19, 21: 20, 22: 21, 23: 22, 24: 23, 25: 24, 27: 25, 28: 26, 31: 27, 32: 28, 33: 29, 34: 30, 35: 31, 36: 32, 37: 33, 38: 34, 39: 35, 40: 36, 41: 37, 42: 38, 43: 39, 44: 40, 46: 41, 47: 42, 48: 43, 49: 44, 50: 45, 51: 46, 52: 47, 53: 48, 54: 49, 55: 50, 56: 51, 57: 52, 58: 53, 59: 54, 60: 55, 61: 56, 62: 57, 63: 58, 64: 59, 65: 60, 67: 61, 70: 62, 72: 63, 73: 64, 74: 65, 75: 66, 76: 67, 77: 68, 78: 69, 79: 70, 80: 71, 81: 72, 82: 73, 84: 74, 85: 75, 86: 76, 87: 77, 88: 78, 89: 79, 90: 80, 92: 81, 93: 82, 94: 83, 95: 84, 96: 85, 97: 86, 98: 87, 99: 88, 100: 89, 101: 90, 102: 91, 103: 92, 104: 93, 105: 94, 106: 95, 107: 96, 108: 97, 109: 98, 110: 99, 111: 100, 112: 101, 113: 102, 114: 103, 115: 104, 116: 105, 117: 106, 118: 107, 119: 108, 120: 109, 121: 110, 122: 111, 123: 112, 124: 113, 125: 114, 126: 115, 127: 116, 128: 117, 129: 118, 130: 119, 131: 120, 132: 121, 133: 122, 134: 123, 135: 124, 136: 125, 137: 126, 138: 127, 139: 128, 140: 129, 141: 130, 142: 131, 143: 132, 144: 133, 145: 134, 146: 135, 147: 136, 148: 137, 149: 138, 150: 139, 151: 140, 152: 141, 153: 142, 154: 143, 155: 144, 156: 145, 157: 146, 158: 147, 159: 148, 160: 149, 161: 150, 162: 151, 163: 152, 164: 153, 165: 154, 166: 155, 167: 156, 168: 157, 169: 158, 170: 159, 171: 160, 172: 161, 173: 162, 174: 163, 175: 164, 176: 165, 177: 166, 178: 167, 179: 168, 180: 169, 181: 170, 182: 171, 12: 0, 26: 0, 29: 0, 30: 0, 45: 0, 66: 0, 68: 0, 69: 0, 71: 0, 83: 0, 91: 0} def relabel(self, labelmap): # label label_copy = labelmap.copy() for k, v in self.id_to_trainid.items(): label_copy[labelmap == k] = v return label_copy.astype(np.uint8) def reduce_one(self, labelmap): return (labelmap - 1).astype(np.uint8) def add_one(self, labelmap): return (labelmap + 1).astype(np.uint8) def evaluate(self, pred_dir, gt_dir): img_cnt = 0 for filename in os.listdir(pred_dir): print(filename) pred_path = os.path.join(pred_dir, filename) gt_path = os.path.join(gt_dir, filename) predmap = ImageHelper.img2np(ImageHelper.read_image(pred_path, tool='pil', mode='P')) gtmap = ImageHelper.img2np(ImageHelper.read_image(gt_path, tool='pil', mode='P')) predmap = self.relabel(predmap) gtmap = self.relabel(gtmap) gtmap[gtmap == 0] = 255 self.seg_running_score.update(predmap[np.newaxis, :, :], gtmap[np.newaxis, :, :]) img_cnt += 1 Log.info('Evaluate {} images'.format(img_cnt)) Log.info('mIOU: {}'.format(self.seg_running_score.get_mean_iou())) Log.info('Pixel ACC: {}'.format(self.seg_running_score.get_pixel_acc()))
def __init__(self, configer): self.configer = configer self.batch_time = AverageMeter() self.data_time = AverageMeter() self.train_losses = AverageMeter() self.val_losses = AverageMeter() self.running_score = RunningScore(configer) self.seg_visualizer = SegVisualizer(configer) self.loss_manager = LossManager(configer) self.module_runner = ModuleRunner(configer) self.model_manager = ModelManager(configer) self.data_loader = DataLoader(configer) self.optim_scheduler = OptimScheduler(configer) self.seg_net = None self.train_loader = None self.val_loader = None self.optimizer = None self.scheduler = None self._init_model()
def __init__(self, configer): self.configer = configer self.seg_running_score = RunningScore(configer) self.id_to_trainid = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 11, 13: 12, 14: 13, 15: 14, 16: 15, 17: 16, 18: 17, 19: 18, 20: 19, 21: 20, 22: 21, 23: 22, 24: 23, 25: 24, 27: 25, 28: 26, 31: 27, 32: 28, 33: 29, 34: 30, 35: 31, 36: 32, 37: 33, 38: 34, 39: 35, 40: 36, 41: 37, 42: 38, 43: 39, 44: 40, 46: 41, 47: 42, 48: 43, 49: 44, 50: 45, 51: 46, 52: 47, 53: 48, 54: 49, 55: 50, 56: 51, 57: 52, 58: 53, 59: 54, 60: 55, 61: 56, 62: 57, 63: 58, 64: 59, 65: 60, 67: 61, 70: 62, 72: 63, 73: 64, 74: 65, 75: 66, 76: 67, 77: 68, 78: 69, 79: 70, 80: 71, 81: 72, 82: 73, 84: 74, 85: 75, 86: 76, 87: 77, 88: 78, 89: 79, 90: 80, 92: 81, 93: 82, 94: 83, 95: 84, 96: 85, 97: 86, 98: 87, 99: 88, 100: 89, 101: 90, 102: 91, 103: 92, 104: 93, 105: 94, 106: 95, 107: 96, 108: 97, 109: 98, 110: 99, 111: 100, 112: 101, 113: 102, 114: 103, 115: 104, 116: 105, 117: 106, 118: 107, 119: 108, 120: 109, 121: 110, 122: 111, 123: 112, 124: 113, 125: 114, 126: 115, 127: 116, 128: 117, 129: 118, 130: 119, 131: 120, 132: 121, 133: 122, 134: 123, 135: 124, 136: 125, 137: 126, 138: 127, 139: 128, 140: 129, 141: 130, 142: 131, 143: 132, 144: 133, 145: 134, 146: 135, 147: 136, 148: 137, 149: 138, 150: 139, 151: 140, 152: 141, 153: 142, 154: 143, 155: 144, 156: 145, 157: 146, 158: 147, 159: 148, 160: 149, 161: 150, 162: 151, 163: 152, 164: 153, 165: 154, 166: 155, 167: 156, 168: 157, 169: 158, 170: 159, 171: 160, 172: 161, 173: 162, 174: 163, 175: 164, 176: 165, 177: 166, 178: 167, 179: 168, 180: 169, 181: 170, 182: 171, 12: 0, 26: 0, 29: 0, 30: 0, 45: 0, 66: 0, 68: 0, 69: 0, 71: 0, 83: 0, 91: 0}
def __init__(self, configer): self.configer = configer self.seg_running_score = RunningScore(configer)
class Trainer(object): """ The class for Pose Estimation. Include train, val, val & 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.running_score = RunningScore(configer) self.seg_visualizer = SegVisualizer(configer) self.loss_manager = LossManager(configer) self.module_runner = ModuleRunner(configer) self.model_manager = ModelManager(configer) self.data_loader = DataLoader(configer) self.optim_scheduler = OptimScheduler(configer) self.seg_net = None self.train_loader = None self.val_loader = None self.optimizer = None self.scheduler = None self._init_model() def _init_model(self): self.seg_net = self.model_manager.semantic_segmentor() self.seg_net = self.module_runner.load_net(self.seg_net) Log.info('Params Group Method: {}'.format( self.configer.get('optim', 'group_method'))) if self.configer.get('optim', 'group_method') == 'decay': params_group = self.group_weight(self.seg_net) else: assert self.configer.get('optim', 'group_method') is None params_group = self._get_parameters() self.optimizer, self.scheduler = self.optim_scheduler.init_optimizer( params_group) self.train_loader = self.data_loader.get_trainloader() self.val_loader = self.data_loader.get_valloader() self.pixel_loss = self.loss_manager.get_seg_loss() @staticmethod def group_weight(module): group_decay = [] group_no_decay = [] for m in module.modules(): if isinstance(m, nn.Linear): group_decay.append(m.weight) if m.bias is not None: group_no_decay.append(m.bias) elif isinstance(m, nn.modules.conv._ConvNd): group_decay.append(m.weight) if m.bias is not None: group_no_decay.append(m.bias) else: if hasattr(m, 'weight'): group_no_decay.append(m.weight) if hasattr(m, 'bias'): group_no_decay.append(m.bias) assert len(list( module.parameters())) == len(group_decay) + len(group_no_decay) groups = [ dict(params=group_decay), dict(params=group_no_decay, weight_decay=.0) ] return groups def _get_parameters(self): bb_lr = [] nbb_lr = [] params_dict = dict(self.seg_net.named_parameters()) for key, value in params_dict.items(): if 'backbone' not in key: nbb_lr.append(value) else: bb_lr.append(value) params = [{ 'params': bb_lr, 'lr': self.configer.get('lr', 'base_lr') }, { 'params': nbb_lr, 'lr': self.configer.get('lr', 'base_lr') * self.configer.get('lr', 'nbb_mult') }] return params def __train(self): """ Train function of every epoch during train phase. """ self.seg_net.train() start_time = time.time() for i, data_dict in enumerate(self.train_loader): if self.configer.get('lr', 'metric') == 'iters': self.scheduler.step(self.configer.get('iters')) else: self.scheduler.step(self.configer.get('epoch')) if self.configer.get('lr', 'is_warm'): self.module_runner.warm_lr(self.configer.get('iters'), self.scheduler, self.optimizer, backbone_list=[ 0, ]) inputs = data_dict['img'] targets = data_dict['labelmap'] self.data_time.update(time.time() - start_time) # Change the data type. # inputs, targets = self.module_runner.to_device(inputs, targets) # Forward pass. outputs = self.seg_net(inputs) # outputs = self.module_utilizer.gather(outputs) # Compute the loss of the train batch & backward. loss = self.pixel_loss(outputs, targets, gathered=self.configer.get( 'network', 'gathered')) if self.configer.exists('train', 'loader') and self.configer.get( 'train', 'loader') == 'ade20k': batch_size = self.configer.get( 'train', 'batch_size') * self.configer.get( 'train', 'batch_per_gpu') self.train_losses.update(loss.item(), batch_size) else: 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.configer.get('epoch'), self.configer.get('iters'), self.configer.get('solver', 'display_iter'), self.module_runner.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('iters') == self.configer.get( 'solver', 'max_iters'): break # Check to val the current model. if self.configer.get('iters') % self.configer.get( 'solver', 'test_interval') == 0: self.__val() self.configer.plus_one('epoch') def __val(self, data_loader=None): """ Validation function during the train phase. """ self.seg_net.eval() start_time = time.time() data_loader = self.val_loader if data_loader is None else data_loader for j, data_dict in enumerate(data_loader): inputs = data_dict['img'] targets = data_dict['labelmap'] with torch.no_grad(): # Change the data type. inputs, targets = self.module_runner.to_device(inputs, targets) # Forward pass. outputs = self.seg_net(inputs) # Compute the loss of the val batch. loss = self.pixel_loss(outputs, targets, gathered=self.configer.get( 'network', 'gathered')) outputs = self.module_runner.gather(outputs) self.val_losses.update(loss.item(), inputs.size(0)) self._update_running_score(outputs[-1], data_dict['meta']) # self.seg_running_score.update(pred.max(1)[1].cpu().numpy(), targets.cpu().numpy()) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.configer.update(['performance'], self.running_score.get_mean_iou()) self.configer.update(['val_loss'], self.val_losses.avg) self.module_runner.save_net(self.seg_net, save_mode='performance') self.module_runner.save_net(self.seg_net, save_mode='val_loss') # 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('Mean IOU: {}\n'.format(self.running_score.get_mean_iou())) Log.info('Pixel ACC: {}\n'.format(self.running_score.get_pixel_acc())) self.batch_time.reset() self.val_losses.reset() self.running_score.reset() self.seg_net.train() def _update_running_score(self, pred, metas): pred = pred.permute(0, 2, 3, 1) for i in range(pred.size(0)): ori_img_size = metas[i]['ori_img_size'] border_size = metas[i]['border_size'] ori_target = metas[i]['ori_target'] total_logits = cv2.resize( pred[i, :border_size[1], :border_size[0]].cpu().numpy(), tuple(ori_img_size), interpolation=cv2.INTER_CUBIC) labelmap = np.argmax(total_logits, axis=-1) self.running_score.update(labelmap[None], ori_target[None]) def train(self): # cudnn.benchmark = True if self.configer.get('network', 'resume') is not None and self.configer.get( 'network', 'resume_val'): self.__val() while self.configer.get('iters') < self.configer.get( 'solver', 'max_iters'): self.__train() self.__val(data_loader=self.data_loader.get_valloader(dataset='val')) self.__val(data_loader=self.data_loader.get_valloader(dataset='train'))