Dev crnn easy

apis/train.py

@@ -12,6 +12,7 @@
 import time
 import os
 from detectron2.data import build_detection_test_loader
+from data import build_lmdb_recognizer_train_loader, build_lmdb_recognizer_test_loader, build_360cc_recognizer_train_loader
 
 from detectron2.engine.defaults import DefaultTrainer
 from detectron2.utils import comm
@@ -24,8 +25,8 @@
 from detectron2.modeling import build_model
 
 
-from data import DatasetMapper, build_detection_train_loader
-from torchtools.optim import RangerLars
+from data import DatasetMapper, build_detection_train_loader, lmdb_dataset
+# from torchtools.optim import RangerLars
 from solver import WarmupCosineAnnealingLR
 from detectron2.solver import build_lr_scheduler, build_optimizer
 from detectron2.solver.build import maybe_add_gradient_clipping
@@ -63,9 +64,13 @@ def __init__(self, cfg):
         # For training, wrap with DDP. But don't need this for inference.
         if comm.get_world_size() > 1:
             model = DistributedDataParallel(
-                model, device_ids=[comm.get_local_rank()], broadcast_buffers=False,find_unused_parameters=True
+                model, device_ids=[comm.get_local_rank()], broadcast_buffers=False, find_unused_parameters=True
             )
-        super(DefaultTrainer, self).__init__(model, data_loader, optimizer)
+
+        self.model = model
+        self.data_loader = data_loader
+        self.optimizer = optimizer
+        super(DefaultTrainer, self).__init__()#model, data_loader, optimizer)
 
         self.scheduler = self.build_lr_scheduler(cfg, optimizer)
         # Assume no other objects need to be checkpointed.
@@ -106,10 +111,16 @@ def build_model(cls, cfg):
 
     @classmethod
     def build_test_loader(cls, cfg, dataset_name):
+        if cfg.DATASETS.TYPE == "CRNN":
+            return build_lmdb_recognizer_test_loader(cfg)
         return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
 
     @classmethod
     def build_train_loader(cls, cfg):
+        if cfg.DATASETS.TYPE == "360CC":
+            return build_360cc_recognizer_train_loader(cfg)
+        elif cfg.DATASETS.TYPE == "lmdb":
+            return build_lmdb_recognizer_train_loader(cfg)
         return build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
 
     @classmethod
@@ -307,8 +318,12 @@ def auto_scale_hyperparams(cfg, data_loader):
         frozen = cfg.is_frozen()
         cfg.defrost()
 
-        iters_per_epoch = len(
-            data_loader.dataset.dataset) // cfg.SOLVER.IMS_PER_BATCH
+        if cfg.DATASETS.TYPE == "360CC":
+            iters_per_epoch = len(data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
+        else:
+            iters_per_epoch = len(
+                data_loader.dataset.dataset) // cfg.SOLVER.IMS_PER_BATCH
+        print("iters_per_epoch:", iters_per_epoch)
         cfg.SOLVER.MAX_ITER *= iters_per_epoch
         cfg.SOLVER.WARMUP_ITERS *= iters_per_epoch
         cfg.SOLVER.WARMUP_FACTOR = 1.0 / cfg.SOLVER.WARMUP_ITERS

configs/config.py

@@ -52,6 +52,16 @@ def add_textnet_config(cfg):
 
     _C.MODEL.DETNET.IMGAUG_PROB = 2.0
 
+    # CRNN config
+    _C.MODEL.CRNN = CN()
+    _C.MODEL.CRNN.NAME = "CRNN"
+    _C.MODEL.CRNN.IMAGE_SIZE_OW = 280
+    _C.MODEL.CRNN.IMAGE_SIZE_H = 32
+    _C.MODEL.CRNN.IMAGE_SIZE_W = 160
+    _C.MODEL.CRNN.NUM_CLASSES = 0
+    _C.MODEL.CRNN.NUM_HIDDEN = 256
+
+
     # rewrite backbone
     _C.MODEL.BACKBONE = CN()
     _C.MODEL.BACKBONE.NAME = "build_resnet"
@@ -81,6 +91,22 @@ def add_textnet_config(cfg):
     _C.INPUT.FORMAT = "RGB"
     _C.INPUT.RESIZE_TYPE = "ResizeShortestEdge"
 
+    # _C.DATASETS = CN()
+    _C.DATASETS.TYPE = "360CC"
+    _C.DATASETS.CHAR_FILE = ""
+    _C.DATASETS.JSON_FILE_TRAIN = ''
+    _C.DATASETS.JSON_FILE_VAL = ''
+    _C.DATASETS.ROOT = ""
+    _C.DATASETS.MEAN = 0.0
+    _C.DATASETS.STD = 0.0
+    _C.DATASETS.ALPHABETS = ''
+
+
+    _C.SOLVER.SHUFFLE = True
+    _C.SOLVER.WORKERS = 4
+    _C.SOLVER.PIN_MEMORY = False
+    _C.SOLVER.OPTIMIZER = ""
+
 
 def add_centernet_config(cfg):
     """

data/__init__.py

@@ -3,4 +3,5 @@
 from .dataset_mapper import DatasetMapper
 from .transforms import *
 from .dataset import *
-from .build import build_detection_train_loader
+from .build import build_detection_train_loader, build_lmdb_recognizer_train_loader, build_lmdb_recognizer_test_loader, build_360cc_recognizer_train_loader
+from .dataset import lmdb_dataset

data/build.py

@@ -21,6 +21,10 @@
 from detectron2.data.detection_utils import check_metadata_consistency
 from detectron2.data.samplers import InferenceSampler, RepeatFactorTrainingSampler, TrainingSampler
 
+from .dataset import lmdb_dataset
+from .dataset.dataset_360cc import Dataset_360CC
+
+
 """
 This file contains the default logic to build a dataloader for training or testing.
 """
@@ -32,6 +36,9 @@
     "get_detection_dataset_dicts",
     "load_proposals_into_dataset",
     "print_instances_class_histogram",
+    "build_lmdb_recognizer_train_loader",
+    "build_lmdb_recognizer_test_loader",
+    "build_360cc_recognizer_train_loader",
 ]
 
 
@@ -339,6 +346,42 @@ def build_detection_test_loader(cfg, dataset_name, mapper=None):
     )
     return data_loader
 
+def build_360cc_recognizer_train_loader(cfg):
+    train_dataset = Dataset_360CC(cfg, is_train=True)
+    train_loader = torch.utils.data.DataLoader(
+        dataset=train_dataset,
+        batch_size=cfg.SOLVER.IMS_PER_BATCH,
+        shuffle=cfg.SOLVER.SHUFFLE,
+        num_workers=cfg.SOLVER.WORKERS,
+        pin_memory=cfg.SOLVER.PIN_MEMORY,
+    )
+
+    return train_loader
+
+def build_lmdb_recognizer_train_loader(cfg):
+    train_dataset = lmdb_dataset.lmdbDataset(root=cfg.DATASETS.TRAIN_ROOT)
+    sampler = None
+    batch_size = cfg.SOLVER.IMS_PER_BATCH
+    if cfg.DATASETS.RANDOM_SAMPLE:
+        sampler = train_dataset.randomSequentialSampler(train_dataset, batch_size)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size,
+        shuffle=True, sampler=sampler,
+        num_workers=int(cfg.SOLVER.WORKERS),
+        collate_fn=lmdb_dataset.alignCollate(imgH=cfg.INPUT.IMG_H, imgW=cfg.INPUT.IMG_W, keep_ratio=cfg.INPUT.KEEP_RATIO))
+
+    return train_loader
+
+def build_lmdb_recognizer_test_loader(cfg):
+    test_dataset = lmdb_dataset.lmdbDataset(
+        root=cfg.DATASETS.TEST_ROOT, transform=lmdb_dataset.resizeNormalize((100, 32)))
+
+    batch_size = cfg.SOLVER.IMS_PER_BATCH
+    test_loader = torch.utils.data.DataLoader(
+        test_dataset, shuffle=True, batch_size=batch_size, num_workers=int(cfg.SOLVER.WORKERS))
+
+    return test_loader
 
 def trivial_batch_collator(batch):
     """

data/dataset/dataset_360cc.py

@@ -0,0 +1,60 @@
+from __future__ import print_function, absolute_import
+import torch.utils.data as data
+import os
+import numpy as np
+import cv2
+
+class Dataset_360CC(data.Dataset):
+    def __init__(self, config, is_train=True):
+
+        self.root = config.DATASETS.ROOT
+        self.is_train = is_train
+        self.inp_h = config.MODEL.CRNN.IMAGE_SIZE_H
+        self.inp_w = config.MODEL.CRNN.IMAGE_SIZE_W
+
+        self.dataset_name = config.DATASETS.TYPE
+
+        self.mean = np.array(config.DATASETS.MEAN, dtype=np.float32)
+        self.std = np.array(config.DATASETS.STD, dtype=np.float32)
+
+        char_file = config.DATASETS.CHAR_FILE
+        with open(char_file, 'rb') as file:
+            char_dict = {num: char.strip().decode('gbk', 'ignore') for num, char in enumerate(file.readlines())}
+
+        txt_file = config.DATASETS.JSON_FILE_TRAIN if is_train else config.DATASETS.JSON_FILE_VAL
+
+        # convert name:indices to name:string
+        self.labels = []
+        with open(txt_file, 'r', encoding='utf-8') as file:
+            contents = file.readlines()
+            for c in contents:
+                imgname = c.split(' ')[0]
+                indices = c.split(' ')[1:]
+                string = ''.join([char_dict[int(idx)] for idx in indices])
+                self.labels.append({imgname: string})
+
+        print("load {} images!".format(self.__len__()))
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+
+        img_name = list(self.labels[idx].keys())[0]
+        img = cv2.imread(os.path.join(self.root, img_name))
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+        img_h, img_w = img.shape
+
+        img = cv2.resize(img, (0,0), fx=self.inp_w / img_w, fy=self.inp_h / img_h, interpolation=cv2.INTER_CUBIC)
+        img = np.reshape(img, (self.inp_h, self.inp_w, 1))
+
+        img = img.astype(np.float32)
+        img = (img/255. - self.mean) / self.std
+        img = img.transpose([2, 0, 1])
+
+        return img, idx
+
+
+
+

data/dataset/dataset_wrapper.py

@@ -0,0 +1,144 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import bisect
+import collections
+import copy
+import math
+from collections import defaultdict
+from torch.utils.data.dataset import ConcatDataset as _ConcatDataset
+
+import numpy as np
+
+class ConcatDataset(_ConcatDataset):
+    """A wrapper of concatenated dataset.
+
+    Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but
+    concat the group flag for image aspect ratio.
+
+    Args:
+        datasets (list[:obj:`Dataset`]): A list of datasets.
+        separate_eval (bool): Whether to evaluate the results
+            separately if it is used as validation dataset.
+            Defaults to True.
+    """
+
+    def __init__(self, datasets, separate_eval=True):
+        super(ConcatDataset, self).__init__(datasets)
+        self.CLASSES = datasets[0].CLASSES
+        self.PALETTE = getattr(datasets[0], 'PALETTE', None)
+        self.separate_eval = separate_eval
+        # if not separate_eval:
+        #     if any([isinstance(ds, CocoDataset) for ds in datasets]):
+        #         raise NotImplementedError(
+        #             'Evaluating concatenated CocoDataset as a whole is not'
+        #             ' supported! Please set "separate_eval=True"')
+        #     elif len(set([type(ds) for ds in datasets])) != 1:
+        #         raise NotImplementedError(
+        #             'All the datasets should have same types')
+
+        if hasattr(datasets[0], 'flag'):
+            flags = []
+            for i in range(0, len(datasets)):
+                flags.append(datasets[i].flag)
+            self.flag = np.concatenate(flags)
+
+    def get_cat_ids(self, idx):
+        """Get category ids of concatenated dataset by index.
+
+        Args:
+            idx (int): Index of data.
+
+        Returns:
+            list[int]: All categories in the image of specified index.
+        """
+
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError(
+                    'absolute value of index should not exceed dataset length')
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        return self.datasets[dataset_idx].get_cat_ids(sample_idx)
+
+    def get_ann_info(self, idx):
+        """Get annotation of concatenated dataset by index.
+
+        Args:
+            idx (int): Index of data.
+
+        Returns:
+            dict: Annotation info of specified index.
+        """
+
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError(
+                    'absolute value of index should not exceed dataset length')
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        return self.datasets[dataset_idx].get_ann_info(sample_idx)
+
+    def evaluate(self, results, logger=None, **kwargs):
+        """Evaluate the results.
+
+        Args:
+            results (list[list | tuple]): Testing results of the dataset.
+            logger (logging.Logger | str | None): Logger used for printing
+                related information during evaluation. Default: None.
+
+        Returns:
+            dict[str: float]: AP results of the total dataset or each separate
+            dataset if `self.separate_eval=True`.
+        """
+        assert len(results) == self.cumulative_sizes[-1], \
+            ('Dataset and results have different sizes: '
+             f'{self.cumulative_sizes[-1]} v.s. {len(results)}')
+
+        # Check whether all the datasets support evaluation
+        for dataset in self.datasets:
+            assert hasattr(dataset, 'evaluate'), \
+                f'{type(dataset)} does not implement evaluate function'
+
+        if self.separate_eval:
+            dataset_idx = -1
+            total_eval_results = dict()
+            for size, dataset in zip(self.cumulative_sizes, self.datasets):
+                start_idx = 0 if dataset_idx == -1 else \
+                    self.cumulative_sizes[dataset_idx]
+                end_idx = self.cumulative_sizes[dataset_idx + 1]
+
+                results_per_dataset = results[start_idx:end_idx]
+                print_log(
+                    f'\nEvaluateing {dataset.ann_file} with '
+                    f'{len(results_per_dataset)} images now',
+                    logger=logger)
+
+                eval_results_per_dataset = dataset.evaluate(
+                    results_per_dataset, logger=logger, **kwargs)
+                dataset_idx += 1
+                for k, v in eval_results_per_dataset.items():
+                    total_eval_results.update({f'{dataset_idx}_{k}': v})
+
+            return total_eval_results
+        elif any([isinstance(ds, CocoDataset) for ds in self.datasets]):
+            raise NotImplementedError(
+                'Evaluating concatenated CocoDataset as a whole is not'
+                ' supported! Please set "separate_eval=True"')
+        elif len(set([type(ds) for ds in self.datasets])) != 1:
+            raise NotImplementedError(
+                'All the datasets should have same types')
+        else:
+            original_data_infos = self.datasets[0].data_infos
+            self.datasets[0].data_infos = sum(
+                [dataset.data_infos for dataset in self.datasets], [])
+            eval_results = self.datasets[0].evaluate(
+                results, logger=logger, **kwargs)
+            self.datasets[0].data_infos = original_data_infos
+            return eval_results