captchaboy/dfff4444

1import logging2import os3import time4 5import cv26import numpy as np7import torch8import yaml9from matplotlib import colors10from matplotlib import pyplot as plt11from torch import Tensor, nn12from torch.utils.data import ConcatDataset13 14class CharsetMapper(object):15    """A simple class to map ids into strings.16 17    It works only when the character set is 1:1 mapping between individual18    characters and individual ids.19    """20 21    def __init__(self,22                 filename='',23                 max_length=30,24                 null_char=u'\u2591'):25        """Creates a lookup table.26 27        Args:28          filename: Path to charset file which maps characters to ids.29          max_sequence_length: The max length of ids and string.30          null_char: A unicode character used to replace '<null>' character.31            the default value is a light shade block '░'.32        """33        self.null_char = null_char34        self.max_length = max_length35 36        self.label_to_char = self._read_charset(filename)37        self.char_to_label = dict(map(reversed, self.label_to_char.items()))38        self.num_classes = len(self.label_to_char)39 40    def _read_charset(self, filename):41        """Reads a charset definition from a tab separated text file.42 43        Args:44          filename: a path to the charset file.45 46        Returns:47          a dictionary with keys equal to character codes and values - unicode48          characters.49        """50        import re51        pattern = re.compile(r'(\d+)\t(.+)')52        charset = {}53        self.null_label = 054        charset[self.null_label] = self.null_char55        with open(filename, 'r') as f:56            for i, line in enumerate(f):57                m = pattern.match(line)58                assert m, f'Incorrect charset file. line #{i}: {line}'59                label = int(m.group(1)) + 160                char = m.group(2)61                charset[label] = char62        return charset63 64    def trim(self, text):65        assert isinstance(text, str)66        return text.replace(self.null_char, '')67 68    def get_text(self, labels, length=None, padding=True, trim=False):69        """ Returns a string corresponding to a sequence of character ids.70        """71        length = length if length else self.max_length72        labels = [l.item() if isinstance(l, Tensor) else int(l) for l in labels]73        if padding:74            labels = labels + [self.null_label] * (length-len(labels))75        text = ''.join([self.label_to_char[label] for label in labels])76        if trim: text = self.trim(text)77        return text78 79    def get_labels(self, text, length=None, padding=True, case_sensitive=False):80        """ Returns the labels of the corresponding text.81        """82        length = length if length else self.max_length83        if padding:84            text = text + self.null_char * (length - len(text))85        if not case_sensitive:86            text = text.lower()87        labels = [self.char_to_label[char] for char in text]88        return labels89 90    def pad_labels(self, labels, length=None):91        length = length if length else self.max_length92 93        return labels + [self.null_label] * (length - len(labels))94 95    @property96    def digits(self):97        return '0123456789'98 99    @property100    def digit_labels(self):101        return self.get_labels(self.digits, padding=False)102 103    @property104    def alphabets(self):105        all_chars = list(self.char_to_label.keys())106        valid_chars = []107        for c in all_chars:108            if c in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ':109                valid_chars.append(c)110        return ''.join(valid_chars)111 112    @property113    def alphabet_labels(self):114        return self.get_labels(self.alphabets, padding=False)115 116 117class Timer(object):118    """A simple timer."""119    def __init__(self):120        self.data_time = 0.121        self.data_diff = 0.122        self.data_total_time = 0.123        self.data_call = 0124        self.running_time = 0.125        self.running_diff = 0.126        self.running_total_time = 0.127        self.running_call = 0128 129    def tic(self):130        self.start_time = time.time()131        self.running_time = self.start_time132 133    def toc_data(self):134        self.data_time = time.time()135        self.data_diff = self.data_time - self.running_time136        self.data_total_time += self.data_diff137        self.data_call += 1138 139    def toc_running(self):140        self.running_time = time.time()141        self.running_diff = self.running_time - self.data_time142        self.running_total_time += self.running_diff143        self.running_call += 1144 145    def total_time(self):146        return self.data_total_time + self.running_total_time147 148    def average_time(self):149        return  self.average_data_time() + self.average_running_time()150 151    def average_data_time(self):152        return self.data_total_time / (self.data_call or 1)153 154    def average_running_time(self):155        return self.running_total_time / (self.running_call or 1)156 157 158class Logger(object):159    _handle = None160    _root = None161 162    @staticmethod163    def init(output_dir, name, phase):164        format = '[%(asctime)s %(filename)s:%(lineno)d %(levelname)s {}] ' \165                '%(message)s'.format(name)166        logging.basicConfig(level=logging.INFO, format=format)167 168        try: os.makedirs(output_dir)169        except: pass170        config_path = os.path.join(output_dir, f'{phase}.txt')171        Logger._handle = logging.FileHandler(config_path)172        Logger._root = logging.getLogger()173 174    @staticmethod175    def enable_file():176        if Logger._handle is None or Logger._root is None:177            raise Exception('Invoke Logger.init() first!')178        Logger._root.addHandler(Logger._handle)179 180    @staticmethod181    def disable_file():182        if Logger._handle is None or Logger._root is None:183            raise Exception('Invoke Logger.init() first!')184        Logger._root.removeHandler(Logger._handle)185 186 187class Config(object):188 189    def __init__(self, config_path, host=True):190        def __dict2attr(d, prefix=''):191            for k, v in d.items():192                if isinstance(v, dict):193                    __dict2attr(v, f'{prefix}{k}_')194                else:195                    if k == 'phase':196                        assert v in ['train', 'test']197                    if k == 'stage':198                        assert v in ['pretrain-vision', 'pretrain-language',199                                     'train-semi-super', 'train-super']200                    self.__setattr__(f'{prefix}{k}', v)201 202        assert os.path.exists(config_path), '%s does not exists!' % config_path203        with open(config_path) as file:204            config_dict = yaml.load(file, Loader=yaml.FullLoader)205        with open('configs/template.yaml') as file:206            default_config_dict = yaml.load(file, Loader=yaml.FullLoader)207        __dict2attr(default_config_dict)208        __dict2attr(config_dict)209        self.global_workdir = os.path.join(self.global_workdir, self.global_name)210 211    def __getattr__(self, item):212        attr = self.__dict__.get(item)213        if attr is None:214            attr = dict()215            prefix = f'{item}_'216            for k, v in self.__dict__.items():217                if k.startswith(prefix):218                    n = k.replace(prefix, '')219                    attr[n] = v220            return attr if len(attr) > 0 else None221        else:222            return attr223 224    def __repr__(self):225        str = 'ModelConfig(\n'226        for i, (k, v) in enumerate(sorted(vars(self).items())):227            str += f'\t({i}): {k} = {v}\n'228        str += ')'229        return str230 231def blend_mask(image, mask, alpha=0.5, cmap='jet', color='b', color_alpha=1.0):232    # normalize mask233    mask = (mask-mask.min()) / (mask.max() - mask.min() + np.finfo(float).eps)234    if mask.shape != image.shape:235        mask = cv2.resize(mask,(image.shape[1], image.shape[0]))236    # get color map237    color_map = plt.get_cmap(cmap)238    mask = color_map(mask)[:,:,:3]239    # convert float to uint8240    mask = (mask * 255).astype(dtype=np.uint8)241 242    # set the basic color243    basic_color = np.array(colors.to_rgb(color)) * 255 244    basic_color = np.tile(basic_color, [image.shape[0], image.shape[1], 1]) 245    basic_color = basic_color.astype(dtype=np.uint8)246    # blend with basic color247    blended_img = cv2.addWeighted(image, color_alpha, basic_color, 1-color_alpha, 0)248    # blend with mask249    blended_img = cv2.addWeighted(blended_img, alpha, mask, 1-alpha, 0)250 251    return blended_img252 253def onehot(label, depth, device=None):254    """ 255    Args:256        label: shape (n1, n2, ..., )257        depth: a scalar258 259    Returns:260        onehot: (n1, n2, ..., depth)261    """262    if not isinstance(label, torch.Tensor):263        label = torch.tensor(label, device=device)264    onehot = torch.zeros(label.size() + torch.Size([depth]), device=device)265    onehot = onehot.scatter_(-1, label.unsqueeze(-1), 1)266 267    return onehot268 269class MyDataParallel(nn.DataParallel):270 271    def gather(self, outputs, target_device):272        r"""273        Gathers tensors from different GPUs on a specified device274        (-1 means the CPU).275        """276        def gather_map(outputs):277            out = outputs[0]278            if isinstance(out, (str, int, float)):279                return out280            if isinstance(out, list) and isinstance(out[0], str):281                return [o for out in outputs for o in out]282            if isinstance(out, torch.Tensor):283                return torch.nn.parallel._functions.Gather.apply(target_device, self.dim, *outputs)284            if out is None:285                return None286            if isinstance(out, dict):287                if not all((len(out) == len(d) for d in outputs)):288                    raise ValueError('All dicts must have the same number of keys')289                return type(out)(((k, gather_map([d[k] for d in outputs]))290                                for k in out))291            return type(out)(map(gather_map, zip(*outputs)))292 293        # Recursive function calls like this create reference cycles.294        # Setting the function to None clears the refcycle.295        try:296            res = gather_map(outputs)297        finally:298            gather_map = None299        return res300 301 302class MyConcatDataset(ConcatDataset):303    def __getattr__(self, k): 304        return getattr(self.datasets[0], k)305