<filename>up/apis/inference.py
import os
import cv2
import numpy as np
import torch
from torch.nn.modules.utils import _pair
from easydict import EasyDict

from up.data.datasets.transforms import build_transformer
from up.data.data_utils import get_image_size
from up.utils.general.cfg_helper import merge_opts_into_cfg
from up.utils.general.log_helper import default_logger as logger
from up.utils.general.registry_factory import MODEL_HELPER_REGISTRY, BATCHING_REGISTRY, IMAGE_READER_REGISTRY
from up.utils.general.registry_factory import INFERENCER_REGISTRY, SAVER_REGISTRY, VISUALIZER_REGISTRY

__all__ = ['BaseInference']


@INFERENCER_REGISTRY.register('base')
class BaseInference(object):
    def __init__(self, config, visualizer, work_dir='./'):
        self.args = config.get('args', {})
        opts = self.args.get('opts', [])
        config = merge_opts_into_cfg(opts, config)
        self.config = config
        self.class_names = visualizer.get('class_names', None)
        self.work_dir = work_dir
        self.vis_dir = self.args['vis_dir']
        self.ckpt = self.args['ckpt']
        self.image_path = self.args['image_path']

        assert self.image_path and os.path.exists(self.image_path), 'Invalid images path.'

        # build DataFetch
        self.build_data()
        logger.info('build data fetcher done')
        # build model
        self.build_model()
        logger.info('build model done')
        # build saver
        self.build_saver()
        logger.info('build saver done')
        # build visualizer
        self.vis_cfg = visualizer
        self.vis_type = self.vis_cfg['type']
        # update vis_dir
        self.vis_cfg['kwargs']['vis_dir'] = self.vis_dir
        self.visualizer = VISUALIZER_REGISTRY.build(self.vis_cfg)
        logger.info('build visualizer done')
        # resume
        self.resume()
        logger.info('load weights done')

    def tensor2numpy(self, x):
        if x is None:
            return x
        if torch.is_tensor(x):
            return x.cpu().numpy()
        if isinstance(x, list):
            x = [_.cpu().numpy() if torch.is_tensor(_) else _ for _ in x]
        return x

    def resume(self):
        checkpoint = self.saver.load_checkpoint(self.ckpt)
        if 'ema' in checkpoint.keys():
            state_dict = checkpoint['model']['ema']['ema_state_dict']
        else:
            state_dict = checkpoint.get('model', checkpoint.get('state_dict', {}))
        self.detector.load(state_dict, strict=False)

    def build_saver(self):
        cfg_saver = self.config['saver']
        _cfg_saver = cfg_saver
        if 'kwargs' not in cfg_saver:
            _cfg_saver = {"type": 'base', "kwargs": {}}
            _cfg_saver['kwargs']['save_cfg'] = cfg_saver
            _cfg_saver['kwargs']['work_dir'] = self.work_dir
        self.saver = SAVER_REGISTRY.build(_cfg_saver)

    def build_model(self):
        model_helper_cfg = self.config['runtime'].get('model_helper', {})
        model_helper_cfg['type'] = model_helper_cfg.get('type', 'base')
        model_helper_cfg['kwargs'] = model_helper_cfg.get('kwargs', {'cfg': self.config['net']})
        self.detector = MODEL_HELPER_REGISTRY.build(model_helper_cfg).cuda().eval()

    def build_data(self):
        data_cfg = self.config['dataset']
        assert 'test' in data_cfg, 'Test dataset config must need !'
        dataset_cfg = data_cfg['test']['dataset']['kwargs']
        self.color_mode = dataset_cfg['image_reader']['kwargs']['color_mode']
        # build image_reader
        self.image_reader = IMAGE_READER_REGISTRY.build(dataset_cfg['image_reader'])

        self.transformer = build_transformer(dataset_cfg['transformer'])
        pad_type = data_cfg['dataloader']['kwargs'].get('pad_type', 'batch_pad')
        pad_value = data_cfg['dataloader']['kwargs'].get('pad_value', 0)
        alignment = data_cfg['dataloader']['kwargs']['alignment']
        self.batch_pad = BATCHING_REGISTRY.get(pad_type)(alignment, pad_value)

    def iterate_image(self, image_dir):
        EXTS = ['jpg', 'jpeg', 'png', 'svg', 'bmp']

        for root, subdirs, subfiles in os.walk(image_dir):
            for filename in subfiles:
                ext = filename.rsplit('.', 1)[-1].lower()
                filepath = os.path.join(root, filename)
                if ext in EXTS:
                    yield filepath

    def map_back(self, output):
        """Map predictions to original image
        Args:
           - output: dict
        Returns:
           - output_list: list of dict,
        """
        origin_images = output['origin_image']
        image_info = output['image_info']
        bboxes = self.tensor2numpy(output['dt_bboxes'])
        batch_size = len(image_info)

        output_list = []
        for b_ix in range(batch_size):

            origin_image = origin_images[b_ix]
            if origin_image.ndim == 3:
                origin_image_h, origin_image_w, _ = origin_image.shape
            else:
                origin_image_h, origin_image_w = origin_image.shape

            img_info = image_info[b_ix]
            unpad_image_h, unpad_image_w = img_info[:2]
            scale_h, scale_w = _pair(img_info[2])
            keep_ix = np.where(bboxes[:, 0] == b_ix)[0]

            # resize bbox
            img_bboxes = bboxes[keep_ix]
            img_bboxes[:, 1] /= scale_w
            img_bboxes[:, 2] /= scale_h
            img_bboxes[:, 3] /= scale_w
            img_bboxes[:, 4] /= scale_h
            img_bboxes = img_bboxes[:, 1:]

            img_output = {
                'image': origin_image,
                'image_info': img_info,
                'dt_bboxes': img_bboxes
            }
            output_list.append(img_output)

        return output_list

    def fetch_single(self, filename):
        img = self.image_reader.read(filename)
        data = EasyDict({
            'filename': filename,
            'origin_image': img,
            'image': img,
            'flipped': False
        })
        data = self.transformer(data)
        scale_factor = data.get('scale_factor', 1)

        image_h, image_w = get_image_size(img)
        new_image_h, new_image_w = get_image_size(data.image)
        pad_w, pad_h = data.get('dw', 0), data.get('dh', 0)
        data.image_info = [new_image_h, new_image_w, scale_factor,
                           image_h, image_w, data.flipped, pad_w, pad_h, filename]
        data.image = data.image.cuda()
        return data

    def fetch(self, filename_list):
        batch = [self.fetch_single(filename) for filename in filename_list]

        batch_keys = list(batch[0].keys())

        def batch_value(key, default=None):
            return [_.get(key, default) for _ in batch]

        data = EasyDict({k: batch_value(k) for k in batch_keys})
        data = self.batch_pad(data)

        return data

    def predict(self):
        if os.path.isdir(self.image_path):
            list_imgs = self.iterate_image(self.image_path)
        else:
            list_imgs = [self.image_path]
        for img_idx, filename in enumerate(list_imgs):
            logger.info('predicting {}:{}'.format(img_idx, filename))
            batch = self.fetch([filename])
            with torch.no_grad():
                output = self.detector(batch)
            output = self.map_back(output)
            self.vis(output)

    def vis(self, outputs):
        for img_idx, output in enumerate(outputs):
            img = output['image']
            if self.color_mode != 'RGB':
                cvt_color_vis = getattr(cv2, 'COLOR_{}2RGB'.format(self.color_mode))
                img = cv2.cvtColor(img, cvt_color_vis)
            boxes = output['dt_bboxes']
            filename = os.path.basename(output['image_info'][-1])
            if self.vis_type == 'plt':
                filename = filename.rsplit('.', 1)[0]
            logger.info('visualizing {}'.format(filename))

            img_h, img_w = img.shape[:2]
            classes = boxes[:, -1].astype(np.int32)
            boxes = boxes[:, :-1]
            self.visualizer.vis(img, boxes, classes, filename)

<filename>up/apis/inference.py
import os
import cv2
import numpy as np
import torch
from torch.nn.modules.utils import _pair
from easydict import EasyDict

from up.data.datasets.transforms import build_transformer
from up.data.data_utils import get_image_size
from up.utils.general.cfg_helper import merge_opts_into_cfg
from up.utils.general.log_helper import default_logger as logger
from up.utils.general.registry_factory import MODEL_HELPER_REGISTRY, BATCHING_REGISTRY, IMAGE_READER_REGISTRY
from up.utils.general.registry_factory import INFERENCER_REGISTRY, SAVER_REGISTRY, VISUALIZER_REGISTRY

__all__ = ['BaseInference']


@INFERENCER_REGISTRY.register('base')
class BaseInference(object):
    def __init__(self, config, visualizer, work_dir='./'):
        self.args = config.get('args', {})
        opts = self.args.get('opts', [])
        config = merge_opts_into_cfg(opts, config)
        self.config = config
        self.class_names = visualizer.get('class_names', None)
        self.work_dir = work_dir
        self.vis_dir = self.args['vis_dir']
        self.ckpt = self.args['ckpt']
        self.image_path = self.args['image_path']

        assert self.image_path and os.path.exists(self.image_path), 'Invalid images path.'

        # build DataFetch
        self.build_data()
        logger.info('build data fetcher done')
        # build model
        self.build_model()
        logger.info('build model done')
        # build saver
        self.build_saver()
        logger.info('build saver done')
        # build visualizer
        self.vis_cfg = visualizer
        self.vis_type = self.vis_cfg['type']
        # update vis_dir
        self.vis_cfg['kwargs']['vis_dir'] = self.vis_dir
        self.visualizer = VISUALIZER_REGISTRY.build(self.vis_cfg)
        logger.info('build visualizer done')
        # resume
        self.resume()
        logger.info('load weights done')

    def tensor2numpy(self, x):
        if x is None:
            return x
        if torch.is_tensor(x):
            return x.cpu().numpy()
        if isinstance(x, list):
            x = [_.cpu().numpy() if torch.is_tensor(_) else _ for _ in x]
        return x

    def resume(self):
        checkpoint = self.saver.load_checkpoint(self.ckpt)
        if 'ema' in checkpoint.keys():
            state_dict = checkpoint['model']['ema']['ema_state_dict']
        else:
            state_dict = checkpoint.get('model', checkpoint.get('state_dict', {}))
        self.detector.load(state_dict, strict=False)

    def build_saver(self):
        cfg_saver = self.config['saver']
        _cfg_saver = cfg_saver
        if 'kwargs' not in cfg_saver:
            _cfg_saver = {"type": 'base', "kwargs": {}}
            _cfg_saver['kwargs']['save_cfg'] = cfg_saver
            _cfg_saver['kwargs']['work_dir'] = self.work_dir
        self.saver = SAVER_REGISTRY.build(_cfg_saver)

    def build_model(self):
        model_helper_cfg = self.config['runtime'].get('model_helper', {})
        model_helper_cfg['type'] = model_helper_cfg.get('type', 'base')
        model_helper_cfg['kwargs'] = model_helper_cfg.get('kwargs', {'cfg': self.config['net']})
        self.detector = MODEL_HELPER_REGISTRY.build(model_helper_cfg).cuda().eval()

    def build_data(self):
        data_cfg = self.config['dataset']
        assert 'test' in data_cfg, 'Test dataset config must need !'
        dataset_cfg = data_cfg['test']['dataset']['kwargs']
        self.color_mode = dataset_cfg['image_reader']['kwargs']['color_mode']
        # build image_reader
        self.image_reader = IMAGE_READER_REGISTRY.build(dataset_cfg['image_reader'])

        self.transformer = build_transformer(dataset_cfg['transformer'])
        pad_type = data_cfg['dataloader']['kwargs'].get('pad_type', 'batch_pad')
        pad_value = data_cfg['dataloader']['kwargs'].get('pad_value', 0)
        alignment = data_cfg['dataloader']['kwargs']['alignment']
        self.batch_pad = BATCHING_REGISTRY.get(pad_type)(alignment, pad_value)

    def iterate_image(self, image_dir):
        EXTS = ['jpg', 'jpeg', 'png', 'svg', 'bmp']

        for root, subdirs, subfiles in os.walk(image_dir):
            for filename in subfiles:
                ext = filename.rsplit('.', 1)[-1].lower()
                filepath = os.path.join(root, filename)
                if ext in EXTS:
                    yield filepath

    def map_back(self, output):
        """Map predictions to original image
        Args:
           - output: dict
        Returns:
           - output_list: list of dict,
        """
        origin_images = output['origin_image']
        image_info = output['image_info']
        bboxes = self.tensor2numpy(output['dt_bboxes'])
        batch_size = len(image_info)

        output_list = []
        for b_ix in range(batch_size):

            origin_image = origin_images[b_ix]
            if origin_image.ndim == 3:
                origin_image_h, origin_image_w, _ = origin_image.shape
            else:
                origin_image_h, origin_image_w = origin_image.shape

            img_info = image_info[b_ix]
            unpad_image_h, unpad_image_w = img_info[:2]
            scale_h, scale_w = _pair(img_info[2])
            keep_ix = np.where(bboxes[:, 0] == b_ix)[0]

            # resize bbox
            img_bboxes = bboxes[keep_ix]
            img_bboxes[:, 1] /= scale_w
            img_bboxes[:, 2] /= scale_h
            img_bboxes[:, 3] /= scale_w
            img_bboxes[:, 4] /= scale_h
            img_bboxes = img_bboxes[:, 1:]

            img_output = {
                'image': origin_image,
                'image_info': img_info,
                'dt_bboxes': img_bboxes
            }
            output_list.append(img_output)

        return output_list

    def fetch_single(self, filename):
        img = self.image_reader.read(filename)
        data = EasyDict({
            'filename': filename,
            'origin_image': img,
            'image': img,
            'flipped': False
        })
        data = self.transformer(data)
        scale_factor = data.get('scale_factor', 1)

        image_h, image_w = get_image_size(img)
        new_image_h, new_image_w = get_image_size(data.image)
        pad_w, pad_h = data.get('dw', 0), data.get('dh', 0)
        data.image_info = [new_image_h, new_image_w, scale_factor,
                           image_h, image_w, data.flipped, pad_w, pad_h, filename]
        data.image = data.image.cuda()
        return data

    def fetch(self, filename_list):
        batch = [self.fetch_single(filename) for filename in filename_list]

        batch_keys = list(batch[0].keys())

        def batch_value(key, default=None):
            return [_.get(key, default) for _ in batch]

        data = EasyDict({k: batch_value(k) for k in batch_keys})
        data = self.batch_pad(data)

        return data

    def predict(self):
        if os.path.isdir(self.image_path):
            list_imgs = self.iterate_image(self.image_path)
        else:
            list_imgs = [self.image_path]
        for img_idx, filename in enumerate(list_imgs):
            logger.info('predicting {}:{}'.format(img_idx, filename))
            batch = self.fetch([filename])
            with torch.no_grad():
                output = self.detector(batch)
            output = self.map_back(output)
            self.vis(output)

    def vis(self, outputs):
        for img_idx, output in enumerate(outputs):
            img = output['image']
            if self.color_mode != 'RGB':
                cvt_color_vis = getattr(cv2, 'COLOR_{}2RGB'.format(self.color_mode))
                img = cv2.cvtColor(img, cvt_color_vis)
            boxes = output['dt_bboxes']
            filename = os.path.basename(output['image_info'][-1])
            if self.vis_type == 'plt':
                filename = filename.rsplit('.', 1)[0]
            logger.info('visualizing {}'.format(filename))

            img_h, img_w = img.shape[:2]
            classes = boxes[:, -1].astype(np.int32)
            boxes = boxes[:, :-1]
            self.visualizer.vis(img, boxes, classes, filename)

# build DataFetch # build model # build saver # build visualizer # update vis_dir # resume # build image_reader Map predictions to original image
         Args:
            - output: dict
         Returns:
            - output_list: list of dict, # resize bbox