添加DBnet代码

requirement.txt

+anyconfig
+tsort
+addict
+natsort
+future==0.18.2
+imgaug==0.4.0
+opencv-python==4.6.0.66
+Polygon3==3.0.8
+pyclipper==1.1.0.post3
+scikit-image==0.16.2
+Shapely==1.6.4.post2

run.sh

+export HIP_VISIBLE_DEVICES=1
+python3 tools/train.py --config_file "config/icdar2015_resnet18_FPN_DBhead_polyLR.yaml"

test/README.MD

+Place the images that you want to detect here. You better named them as such:
+img_10.jpg
+img_11.jpg
+img_{img_id}.jpg
+
+For predicting single images, you can change the `img_path` in the `/tools/predict.py` to your image number.
+
+The result will be saved in the output_folder(default is test/output) you give in predict.sh
\ No newline at end of file

tools/__init__.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/12/8 13:14
+# @Author  : zhoujun
\ No newline at end of file

tools/eval.py

+# -*- coding: utf-8 -*-
+# @Time    : 2018/6/11 15:54
+# @Author  : zhoujun
+import os
+import sys
+import pathlib
+__dir__ = pathlib.Path(os.path.abspath(__file__))
+sys.path.append(str(__dir__))
+sys.path.append(str(__dir__.parent.parent))
+# project = 'DBNet.pytorch'  # 工作项目根目录
+# sys.path.append(os.getcwd().split(project)[0] + project)
+
+import argparse
+import time
+import torch
+from tqdm.auto import tqdm
+
+
+class EVAL():
+    def __init__(self, model_path, gpu_id=0):
+        from models import build_model
+        from data_loader import get_dataloader
+        from post_processing import get_post_processing
+        from utils import get_metric
+        self.gpu_id = gpu_id
+        if self.gpu_id is not None and isinstance(self.gpu_id, int) and torch.cuda.is_available():
+            self.device = torch.device("cuda:%s" % self.gpu_id)
+            torch.backends.cudnn.benchmark = True
+        else:
+            self.device = torch.device("cpu")
+        checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+        config = checkpoint['config']
+        config['arch']['backbone']['pretrained'] = False
+
+        self.validate_loader = get_dataloader(config['dataset']['validate'], config['distributed'])
+
+        self.model = build_model(config['arch'])
+        self.model.load_state_dict(checkpoint['state_dict'])
+        self.model.to(self.device)
+
+        self.post_process = get_post_processing(config['post_processing'])
+        self.metric_cls = get_metric(config['metric'])
+
+    def eval(self):
+        self.model.eval()
+        # torch.cuda.empty_cache()  # speed up evaluating after training finished
+        raw_metrics = []
+        total_frame = 0.0
+        total_time = 0.0
+        for i, batch in tqdm(enumerate(self.validate_loader), total=len(self.validate_loader), desc='test model'):
+            with torch.no_grad():
+                # 数据进行转换和丢到gpu
+                for key, value in batch.items():
+                    if value is not None:
+                        if isinstance(value, torch.Tensor):
+                            batch[key] = value.to(self.device)
+                start = time.time()
+                preds = self.model(batch['img'])
+                boxes, scores = self.post_process(batch, preds,is_output_polygon=self.metric_cls.is_output_polygon)
+                total_frame += batch['img'].size()[0]
+                total_time += time.time() - start
+                raw_metric = self.metric_cls.validate_measure(batch, (boxes, scores))
+                raw_metrics.append(raw_metric)
+        metrics = self.metric_cls.gather_measure(raw_metrics)
+        print('FPS:{}'.format(total_frame / total_time))
+        return metrics['recall'].avg, metrics['precision'].avg, metrics['fmeasure'].avg
+
+
+def init_args():
+    parser = argparse.ArgumentParser(description='DBNet.pytorch')
+    parser.add_argument('--model_path', required=False,default='output/DBNet_resnet18_FPN_DBHead/checkpoint/1.pth', type=str)
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = init_args()
+    eval = EVAL(args.model_path)
+    result = eval.eval()
+    print(result)

tools/predict.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/24 12:06
+# @Author  : zhoujun
+
+import os
+import sys
+import pathlib
+__dir__ = pathlib.Path(os.path.abspath(__file__))
+sys.path.append(str(__dir__))
+sys.path.append(str(__dir__.parent.parent))
+
+# project = 'DBNet.pytorch'  # 工作项目根目录
+# sys.path.append(os.getcwd().split(project)[0] + project)
+import time
+import cv2
+import torch
+
+from data_loader import get_transforms
+from models import build_model
+from post_processing import get_post_processing
+
+
+def resize_image(img, short_size):
+    height, width, _ = img.shape
+    if height < width:
+        new_height = short_size
+        new_width = new_height / height * width
+    else:
+        new_width = short_size
+        new_height = new_width / width * height
+    new_height = int(round(new_height / 32) * 32)
+    new_width = int(round(new_width / 32) * 32)
+    resized_img = cv2.resize(img, (new_width, new_height))
+    return resized_img
+
+
+class Pytorch_model:
+    def __init__(self, model_path, post_p_thre=0.7, gpu_id=None):
+        '''
+        初始化pytorch模型
+        :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件)
+        :param gpu_id: 在哪一块gpu上运行
+        '''
+        self.gpu_id = gpu_id
+
+        if self.gpu_id is not None and isinstance(self.gpu_id, int) and torch.cuda.is_available():
+            self.device = torch.device("cuda:%s" % self.gpu_id)
+        else:
+            self.device = torch.device("cpu")
+        print('device:', self.device)
+        checkpoint = torch.load(model_path, map_location=self.device)
+
+        config = checkpoint['config']
+        config['arch']['backbone']['pretrained'] = False
+        self.model = build_model(config['arch'])
+        self.post_process = get_post_processing(config['post_processing'])
+        self.post_process.box_thresh = post_p_thre
+        self.img_mode = config['dataset']['train']['dataset']['args']['img_mode']
+        self.model.load_state_dict(checkpoint['state_dict'])
+        self.model.to(self.device)
+        self.model.eval()
+
+        self.transform = []
+        for t in config['dataset']['train']['dataset']['args']['transforms']:
+            if t['type'] in ['ToTensor', 'Normalize']:
+                self.transform.append(t)
+        self.transform = get_transforms(self.transform)
+
+    def predict(self, img_path: str, is_output_polygon=False, short_size: int = 1024):
+        '''
+        对传入的图像进行预测，支持图像地址,opecv 读取图片，偏慢
+        :param img_path: 图像地址
+        :param is_numpy:
+        :return:
+        '''
+        assert os.path.exists(img_path), 'file is not exists'
+        img = cv2.imread(img_path, 1 if self.img_mode != 'GRAY' else 0)
+        if self.img_mode == 'RGB':
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        h, w = img.shape[:2]
+        img = resize_image(img, short_size)
+        # 将图片由(w,h)变为(1,img_channel,h,w)
+        tensor = self.transform(img)
+        tensor = tensor.unsqueeze_(0)
+
+        tensor = tensor.to(self.device)
+        batch = {'shape': [(h, w)]}
+        with torch.no_grad():
+            if str(self.device).__contains__('cuda'):
+                torch.cuda.synchronize(self.device)
+            start = time.time()
+            preds = self.model(tensor)
+            if str(self.device).__contains__('cuda'):
+                torch.cuda.synchronize(self.device)
+            box_list, score_list = self.post_process(batch, preds, is_output_polygon=is_output_polygon)
+            box_list, score_list = box_list[0], score_list[0]
+            if len(box_list) > 0:
+                if is_output_polygon:
+                    idx = [x.sum() > 0 for x in box_list]
+                    box_list = [box_list[i] for i, v in enumerate(idx) if v]
+                    score_list = [score_list[i] for i, v in enumerate(idx) if v]
+                else:
+                    idx = box_list.reshape(box_list.shape[0], -1).sum(axis=1) > 0  # 去掉全为0的框
+                    box_list, score_list = box_list[idx], score_list[idx]
+            else:
+                box_list, score_list = [], []
+            t = time.time() - start
+        return preds[0, 0, :, :].detach().cpu().numpy(), box_list, score_list, t
+
+
+def save_depoly(model, input, save_path):
+    traced_script_model = torch.jit.trace(model, input)
+    traced_script_model.save(save_path)
+
+
+def init_args():
+    import argparse
+    parser = argparse.ArgumentParser(description='DBNet.pytorch')
+    parser.add_argument('--model_path', default=r'model_best.pth', type=str)
+    parser.add_argument('--input_folder', default='./test/input', type=str, help='img path for predict')
+    parser.add_argument('--output_folder', default='./test/output', type=str, help='img path for output')
+    parser.add_argument('--thre', default=0.3,type=float, help='the thresh of post_processing')
+    parser.add_argument('--polygon', action='store_true', help='output polygon or box')
+    parser.add_argument('--show', action='store_true', help='show result')
+    parser.add_argument('--save_resut', action='store_true', help='save box and score to txt file')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    import pathlib
+    from tqdm import tqdm
+    import matplotlib.pyplot as plt
+    from utils.util import show_img, draw_bbox, save_result, get_file_list
+
+    args = init_args()
+    print(args)
+    os.environ['CUDA_VISIBLE_DEVICES'] = str('0')
+    # 初始化网络
+    model = Pytorch_model(args.model_path, post_p_thre=args.thre, gpu_id=0)
+    img_folder = pathlib.Path(args.input_folder)
+    for img_path in tqdm(get_file_list(args.input_folder, p_postfix=['.jpg'])):
+        preds, boxes_list, score_list, t = model.predict(img_path, is_output_polygon=args.polygon)
+        img = draw_bbox(cv2.imread(img_path)[:, :, ::-1], boxes_list)
+        if args.show:
+            show_img(preds)
+            show_img(img, title=os.path.basename(img_path))
+            plt.show()
+        # 保存结果到路径
+        os.makedirs(args.output_folder, exist_ok=True)
+        img_path = pathlib.Path(img_path)
+        output_path = os.path.join(args.output_folder, img_path.stem + '_result.jpg')
+        pred_path = os.path.join(args.output_folder, img_path.stem + '_pred.jpg')
+        cv2.imwrite(output_path, img[:, :, ::-1])
+        cv2.imwrite(pred_path, preds * 255)
+        save_result(output_path.replace('_result.jpg', '.txt'), boxes_list, score_list, args.polygon)

tools/train.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/23 22:00
+# @Author  : zhoujun
+
+from __future__ import print_function
+
+import argparse
+import os
+
+import anyconfig
+
+
+def init_args():
+    parser = argparse.ArgumentParser(description='DBNet.pytorch')
+    parser.add_argument('--config_file', default='config/open_dataset_resnet18_FPN_DBhead_polyLR.yaml', type=str)
+    parser.add_argument('--local_rank', dest='local_rank', default=0, type=int, help='Use distributed training')
+
+    args = parser.parse_args()
+    return args
+
+
+def main(config):
+    import torch
+    from models import build_model, build_loss
+    from data_loader import get_dataloader
+    from trainer import Trainer
+    from post_processing import get_post_processing
+    from utils import get_metric
+    if torch.cuda.device_count() > 1:
+        torch.cuda.set_device(args.local_rank)
+        torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=torch.cuda.device_count(), rank=args.local_rank)
+        config['distributed'] = True
+    else:
+        config['distributed'] = False
+    config['local_rank'] = args.local_rank
+
+    train_loader = get_dataloader(config['dataset']['train'], config['distributed'])
+    assert train_loader is not None
+    if 'validate' in config['dataset']:
+        validate_loader = get_dataloader(config['dataset']['validate'], False)
+    else:
+        validate_loader = None
+
+    criterion = build_loss(config['loss']).cuda()
+
+    config['arch']['backbone']['in_channels'] = 3 if config['dataset']['train']['dataset']['args']['img_mode'] != 'GRAY' else 1
+    model = build_model(config['arch'])
+
+    post_p = get_post_processing(config['post_processing'])
+    metric = get_metric(config['metric'])
+
+    trainer = Trainer(config=config,
+                      model=model,
+                      criterion=criterion,
+                      train_loader=train_loader,
+                      post_process=post_p,
+                      metric_cls=metric,
+                      validate_loader=validate_loader)
+    trainer.train()
+
+
+if __name__ == '__main__':
+    import sys
+    import pathlib
+    __dir__ = pathlib.Path(os.path.abspath(__file__))
+    sys.path.append(str(__dir__))
+    sys.path.append(str(__dir__.parent.parent))
+    # project = 'DBNet.pytorch'  # 工作项目根目录
+    # sys.path.append(os.getcwd().split(project)[0] + project)
+
+    from utils import parse_config
+
+    args = init_args()
+    assert os.path.exists(args.config_file)
+    config = anyconfig.load(open(args.config_file, 'rb'))
+    if 'base' in config:
+        config = parse_config(config)
+    main(config)

trainer/__init__.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/23 21:58
+# @Author  : zhoujun
+from .trainer import Trainer
\ No newline at end of file

trainer/trainer.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/23 21:58
+# @Author  : zhoujun
+import time
+
+import torch
+import torchvision.utils as vutils
+from tqdm import tqdm
+
+from base import BaseTrainer
+from utils import WarmupPolyLR, runningScore, cal_text_score
+
+
+class Trainer(BaseTrainer):
+    def __init__(self, config, model, criterion, train_loader, validate_loader, metric_cls, post_process=None):
+        super(Trainer, self).__init__(config, model, criterion)
+        self.show_images_iter = self.config['trainer']['show_images_iter']
+        self.train_loader = train_loader
+        if validate_loader is not None:
+            assert post_process is not None and metric_cls is not None
+        self.validate_loader = validate_loader
+        self.post_process = post_process
+        self.metric_cls = metric_cls
+        self.train_loader_len = len(train_loader)
+        if self.config['lr_scheduler']['type'] == 'WarmupPolyLR':
+            warmup_iters = config['lr_scheduler']['args']['warmup_epoch'] * self.train_loader_len
+            if self.start_epoch > 1:
+                self.config['lr_scheduler']['args']['last_epoch'] = (self.start_epoch - 1) * self.train_loader_len
+            self.scheduler = WarmupPolyLR(self.optimizer, max_iters=self.epochs * self.train_loader_len,
+                                          warmup_iters=warmup_iters, **config['lr_scheduler']['args'])
+        if self.validate_loader is not None:
+            self.logger_info(
+                'train dataset has {} samples,{} in dataloader, validate dataset has {} samples,{} in dataloader'.format(
+                    len(self.train_loader.dataset), self.train_loader_len, len(self.validate_loader.dataset), len(self.validate_loader)))
+        else:
+            self.logger_info('train dataset has {} samples,{} in dataloader'.format(len(self.train_loader.dataset), self.train_loader_len))
+
+    def _train_epoch(self, epoch):
+        self.model.train()
+        epoch_start = time.time()
+        batch_start = time.time()
+        train_loss = 0.
+        running_metric_text = runningScore(2)
+        lr = self.optimizer.param_groups[0]['lr']
+
+        for i, batch in enumerate(self.train_loader):
+            if i >= self.train_loader_len:
+                break
+            self.global_step += 1
+            lr = self.optimizer.param_groups[0]['lr']
+
+            # 数据进行转换和丢到gpu
+            for key, value in batch.items():
+                if value is not None:
+                    if isinstance(value, torch.Tensor):
+                        batch[key] = value.to(self.device)
+            cur_batch_size = batch['img'].size()[0]
+
+            preds = self.model(batch['img'])
+            loss_dict = self.criterion(preds, batch)
+            # backward
+            self.optimizer.zero_grad()
+            loss_dict['loss'].backward()
+            self.optimizer.step()
+            if self.config['lr_scheduler']['type'] == 'WarmupPolyLR':
+                self.scheduler.step()
+            # acc iou
+            score_shrink_map = cal_text_score(preds[:, 0, :, :], batch['shrink_map'], batch['shrink_mask'], running_metric_text,
+                                              thred=self.config['post_processing']['args']['thresh'])
+
+            # loss 和 acc 记录到日志
+            loss_str = 'loss: {:.4f}, '.format(loss_dict['loss'].item())
+            for idx, (key, value) in enumerate(loss_dict.items()):
+                loss_dict[key] = value.item()
+                if key == 'loss':
+                    continue
+                loss_str += '{}: {:.4f}'.format(key, loss_dict[key])
+                if idx < len(loss_dict) - 1:
+                    loss_str += ', '
+
+            train_loss += loss_dict['loss']
+            acc = score_shrink_map['Mean Acc']
+            iou_shrink_map = score_shrink_map['Mean IoU']
+
+            if self.global_step % self.log_iter == 0:
+                batch_time = time.time() - batch_start
+                self.logger_info(
+                    '[{}/{}], [{}/{}], global_step: {}, speed: {:.1f} samples/sec, acc: {:.4f}, iou_shrink_map: {:.4f}, {}, lr:{:.6}, time:{:.2f}'.format(
+                        epoch, self.epochs, i + 1, self.train_loader_len, self.global_step, self.log_iter * cur_batch_size / batch_time, acc,
+                        iou_shrink_map, loss_str, lr, batch_time))
+                batch_start = time.time()
+
+            if self.tensorboard_enable and self.config['local_rank'] == 0:
+                # write tensorboard
+                for key, value in loss_dict.items():
+                    self.writer.add_scalar('TRAIN/LOSS/{}'.format(key), value, self.global_step)
+                self.writer.add_scalar('TRAIN/ACC_IOU/acc', acc, self.global_step)
+                self.writer.add_scalar('TRAIN/ACC_IOU/iou_shrink_map', iou_shrink_map, self.global_step)
+                self.writer.add_scalar('TRAIN/lr', lr, self.global_step)
+                if self.global_step % self.show_images_iter == 0:
+                    # show images on tensorboard
+                    self.inverse_normalize(batch['img'])
+                    self.writer.add_images('TRAIN/imgs', batch['img'], self.global_step)
+                    # shrink_labels and threshold_labels
+                    shrink_labels = batch['shrink_map']
+                    threshold_labels = batch['threshold_map']
+                    shrink_labels[shrink_labels <= 0.5] = 0
+                    shrink_labels[shrink_labels > 0.5] = 1
+                    show_label = torch.cat([shrink_labels, threshold_labels])
+                    show_label = vutils.make_grid(show_label.unsqueeze(1), nrow=cur_batch_size, normalize=False, padding=20, pad_value=1)
+                    self.writer.add_image('TRAIN/gt', show_label, self.global_step)
+                    # model output
+                    show_pred = []
+                    for kk in range(preds.shape[1]):
+                        show_pred.append(preds[:, kk, :, :])
+                    show_pred = torch.cat(show_pred)
+                    show_pred = vutils.make_grid(show_pred.unsqueeze(1), nrow=cur_batch_size, normalize=False, padding=20, pad_value=1)
+                    self.writer.add_image('TRAIN/preds', show_pred, self.global_step)
+        return {'train_loss': train_loss / self.train_loader_len, 'lr': lr, 'time': time.time() - epoch_start,
+                'epoch': epoch}
+
+    def _eval(self, epoch):
+        self.model.eval()
+        # torch.cuda.empty_cache()  # speed up evaluating after training finished
+        raw_metrics = []
+        total_frame = 0.0
+        total_time = 0.0
+        for i, batch in tqdm(enumerate(self.validate_loader), total=len(self.validate_loader), desc='test model'):
+            with torch.no_grad():
+                # 数据进行转换和丢到gpu
+                for key, value in batch.items():
+                    if value is not None:
+                        if isinstance(value, torch.Tensor):
+                            batch[key] = value.to(self.device)
+                start = time.time()
+                preds = self.model(batch['img'])
+                boxes, scores = self.post_process(batch, preds,is_output_polygon=self.metric_cls.is_output_polygon)
+                total_frame += batch['img'].size()[0]
+                total_time += time.time() - start
+                raw_metric = self.metric_cls.validate_measure(batch, (boxes, scores))
+                raw_metrics.append(raw_metric)
+        metrics = self.metric_cls.gather_measure(raw_metrics)
+        self.logger_info('FPS:{}'.format(total_frame / total_time))
+        return metrics['recall'].avg, metrics['precision'].avg, metrics['fmeasure'].avg
+
+    def _on_epoch_finish(self):
+        self.logger_info('[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
+            self.epoch_result['epoch'], self.epochs, self.epoch_result['train_loss'], self.epoch_result['time'],
+            self.epoch_result['lr']))
+        net_save_path = '{}/model_latest.pth'.format(self.checkpoint_dir)
+        net_save_path_best = '{}/model_best.pth'.format(self.checkpoint_dir)
+
+        if self.config['local_rank'] == 0:
+            self._save_checkpoint(self.epoch_result['epoch'], net_save_path)
+            save_best = False
+            if self.validate_loader is not None and self.metric_cls is not None:  # 使用f1作为最优模型指标
+                recall, precision, hmean = self._eval(self.epoch_result['epoch'])
+
+                if self.tensorboard_enable:
+                    self.writer.add_scalar('EVAL/recall', recall, self.global_step)
+                    self.writer.add_scalar('EVAL/precision', precision, self.global_step)
+                    self.writer.add_scalar('EVAL/hmean', hmean, self.global_step)
+                self.logger_info('test: recall: {:.6f}, precision: {:.6f}, f1: {:.6f}'.format(recall, precision, hmean))
+
+                if hmean >= self.metrics['hmean']:
+                    save_best = True
+                    self.metrics['train_loss'] = self.epoch_result['train_loss']
+                    self.metrics['hmean'] = hmean
+                    self.metrics['precision'] = precision
+                    self.metrics['recall'] = recall
+                    self.metrics['best_model_epoch'] = self.epoch_result['epoch']
+            else:
+                if self.epoch_result['train_loss'] <= self.metrics['train_loss']:
+                    save_best = True
+                    self.metrics['train_loss'] = self.epoch_result['train_loss']
+                    self.metrics['best_model_epoch'] = self.epoch_result['epoch']
+            best_str = 'current best, '
+            for k, v in self.metrics.items():
+                best_str += '{}: {:.6f}, '.format(k, v)
+            self.logger_info(best_str)
+            if save_best:
+                import shutil
+                shutil.copy(net_save_path, net_save_path_best)
+                self.logger_info("Saving current best: {}".format(net_save_path_best))
+            else:
+                self.logger_info("Saving checkpoint: {}".format(net_save_path))
+
+
+    def _on_train_finish(self):
+        for k, v in self.metrics.items():
+            self.logger_info('{}:{}'.format(k, v))
+        self.logger_info('finish train')

utils/__init__.py

+# -*- coding: utf-8 -*-
+# @Time    : 2019/8/23 21:58
+# @Author  : zhoujun
+from .util import *
+from .metrics import *
+from .schedulers import *
+from .cal_recall.script import  cal_recall_precison_f1
+from .ocr_metric import get_metric
\ No newline at end of file