# Copyright (c) Alibaba, Inc. and its affiliates.
import json
import logging
import os
import pickle
import cv2
import mmcv
import numpy as np
import torch
from mmcv.parallel import collate, scatter_kwargs
from PIL import Image, ImageFile
from torch.hub import load_state_dict_from_url
from torchvision.transforms import Compose
from easycv.datasets.registry import PIPELINES
from easycv.file import io
from easycv.file.utils import is_url_path
from easycv.framework.errors import ValueError
from easycv.models.builder import build_model
from easycv.utils.checkpoint import load_checkpoint
from easycv.utils.config_tools import Config, mmcv_config_fromfile
from easycv.utils.constant import CACHE_DIR
from easycv.utils.logger import get_root_logger
from easycv.utils.mmlab_utils import (dynamic_adapt_for_mmlab,
remove_adapt_for_mmlab)
from easycv.utils.registry import build_from_cfg
[docs]class NumpyToPIL(object):
def __call__(self, results):
img = results['img']
results['img'] = Image.fromarray(np.uint8(img)).convert('RGB')
return results
[docs]class Predictor(object):
[docs] def __init__(self, model_path, numpy_to_pil=True):
self.model_path = model_path
self.numpy_to_pil = numpy_to_pil
assert io.exists(self.model_path), f'{self.model_path} does not exists'
with io.open(self.model_path, 'rb') as infile:
checkpoint = torch.load(infile, map_location='cpu')
assert 'meta' in checkpoint and 'config' in checkpoint[
'meta'], 'meta.config is missing from checkpoint'
config_str = checkpoint['meta']['config']
# get config
basename = os.path.basename(self.model_path)
fname, _ = os.path.splitext(basename)
self.local_config_file = os.path.join(CACHE_DIR,
f'{fname}_config.json')
if not os.path.exists(CACHE_DIR):
os.makedirs(CACHE_DIR)
with open(self.local_config_file, 'w') as ofile:
ofile.write(config_str)
self.cfg = mmcv_config_fromfile(self.local_config_file)
# build model
self.model = build_model(self.cfg.model)
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
map_location = 'cpu' if self.device == 'cpu' else 'cuda'
self.ckpt = load_checkpoint(
self.model, self.model_path, map_location=map_location)
self.model.to(self.device)
self.model.eval()
# build pipeline
pipeline = [
build_from_cfg(p, PIPELINES) for p in self.cfg.test_pipeline
]
if self.numpy_to_pil:
pipeline = [NumpyToPIL()] + pipeline
self.pipeline = Compose(pipeline)
[docs] def preprocess(self, image_list):
# only perform transform to img
output_imgs_list = []
for img in image_list:
tmp_input = {'img': img}
tmp_results = self.pipeline(tmp_input)
output_imgs_list.append(tmp_results['img'])
return output_imgs_list
[docs] def predict_batch(self, image_batch, **forward_kwargs):
""" predict using batched data
Args:
image_batch(torch.Tensor): tensor with shape [N, 3, H, W]
forward_kwargs: kwargs for additional parameters
Return:
output: the output of model.forward, list or tuple
"""
with torch.no_grad():
output = self.model.forward(
image_batch.to(self.device), **forward_kwargs)
return output
[docs]class OutputProcessor(object):
"""Base output processor for processing model outputs.
"""
[docs] def __init__(self):
pass
def _get_batch_size(self, inputs):
for k, batch_v in inputs.items():
if isinstance(batch_v, dict):
batch_size = self._get_batch_size(batch_v)
elif batch_v is not None:
batch_size = len(batch_v)
break
else:
batch_size = 1
return batch_size
def _extract_ith_result(self, inputs, i, out_i):
for k, batch_v in inputs.items():
if isinstance(batch_v, dict):
out_i[k] = {}
self._extract_ith_result(batch_v, i, out_i[k])
elif batch_v is not None:
out_i[k] = batch_v[i]
else:
out_i[k] = None
return out_i
[docs] def process_single(self, inputs):
"""Process outputs of single sample.
If you need add some processing ops, you need to reimplement it.
"""
return inputs
def __call__(self, inputs):
"""Process model batch outputs.
The "inputs" should be dict format as follows:
{
"key1": torch.Tensor or list, the first dimension should be batch_size,
"key2": torch.Tensor or list, the first dimension should be batch_size,
...
}
"""
outputs = []
batch_size = self._get_batch_size(inputs)
for i in range(batch_size):
out_i = self._extract_ith_result(inputs, i, {})
out_i = self.process_single(out_i)
outputs.append(out_i)
return outputs
[docs]class PredictorV2(object):
"""Base predict pipeline.
Args:
model_path (str): Path of model path.
config_file (Optinal[str]): config file path for model and processor to init. Defaults to None.
batch_size (int): batch size for forward.
device (str | torch.device): Support str('cuda' or 'cpu') or torch.device, if is None, detect device automatically.
save_results (bool): Whether to save predict results.
save_path (str): File path for saving results, only valid when `save_results` is True.
pipelines (list[dict]): Data pipeline configs.
input_processor_threads (int): Number of processes to process inputs.
mode (str): The image mode into the model.
"""
[docs] def __init__(self,
model_path,
config_file=None,
batch_size=1,
device=None,
save_results=False,
save_path=None,
pipelines=None,
input_processor_threads=8,
mode='BGR'):
self.logger = get_root_logger()
self.model_path = model_path
self.batch_size = batch_size
self.save_results = save_results
self.save_path = save_path
self.config_file = config_file
self.pipelines = pipelines
self.input_processor_threads = input_processor_threads
self.mode = mode
if self.save_results:
assert self.save_path is not None
self.device = device
if self.device is None:
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
if config_file is not None:
if isinstance(config_file, str):
self.cfg = mmcv_config_fromfile(config_file)
else:
self.cfg = config_file
else:
self.cfg = self._load_cfg_from_ckpt(self.model_path)
if self.cfg is None:
raise ValueError('Please provide "config_file"!')
if self.cfg.get('predict', None) is not None:
self._sync_cfg_predict(self.cfg.predict)
# avoid unnecessarily loading backbone weights from url
if 'model' in self.cfg and 'pretrained' in self.cfg.model:
self.cfg.model.pretrained = None
self.model = self.prepare_model()
self.input_processor = None
self.output_processor = None
[docs] def get_output_processor(self):
return OutputProcessor()
def _sync_cfg_predict(self, predict_cfg):
if predict_cfg.get('type', None) is not None:
assert predict_cfg[
'type'] == self.__class__.__name__, f'Predictor name is not equal {predict_cfg["type"]} != {self.__class__.__name__}'
for k, v in predict_cfg.items():
if k == 'type':
continue
setattr(self, k, v)
self.logger.warning(
f'Set "{self.__class__.__name__}.{k}" to "{v}" !')
def _load_cfg_from_ckpt(self, model_path):
if is_url_path(model_path):
ckpt = load_state_dict_from_url(model_path)
else:
with io.open(model_path, 'rb') as infile:
ckpt = torch.load(infile, map_location='cpu')
cfg = None
if 'meta' in ckpt and 'config' in ckpt['meta']:
cfg = ckpt['meta']['config']
if isinstance(cfg, dict):
cfg = Config(cfg)
elif isinstance(cfg, str):
cfg = Config(json.loads(cfg))
return cfg
[docs] def prepare_model(self):
"""Build model from config file by default.
If the model is not loaded from a configuration file, e.g. torch jit model, you need to reimplement it.
"""
model = self._build_model()
model.to(self.device)
model.eval()
load_checkpoint(model, self.model_path, map_location='cpu')
return model
def _build_model(self):
# Use mmdet model
dynamic_adapt_for_mmlab(self.cfg)
model = build_model(self.cfg.model)
# remove adapt for mmdet to avoid conflict using mmdet models
remove_adapt_for_mmlab(self.cfg)
return model
[docs] def model_forward(self, inputs):
"""Model forward.
If you need refactor model forward, you need to reimplement it.
"""
with torch.no_grad():
outputs = self.model(**inputs, mode='test')
return outputs
def _to_device(self, inputs):
target_gpus = [-1] if str(
self.device) == 'cpu' else [torch.cuda.current_device()]
_, kwargs = scatter_kwargs(None, inputs, target_gpus=target_gpus)
return kwargs[0]
[docs] def dump(self, obj, save_path, mode='wb'):
with open(save_path, mode) as f:
f.write(pickle.dumps(obj))
def __call__(self, inputs, keep_inputs=False):
if self.input_processor is None:
self.input_processor = self.get_input_processor()
if self.output_processor is None:
self.output_processor = self.get_output_processor()
# TODO: fault tolerance
if isinstance(inputs, (str, np.ndarray, ImageFile.ImageFile)):
inputs = [inputs]
results_list = []
prog_bar = mmcv.ProgressBar(len(inputs))
for i in range(0, len(inputs), self.batch_size):
batch_inputs = inputs[i:min(len(inputs), i + self.batch_size)]
batch_outputs = self.input_processor(batch_inputs)
if len(batch_outputs) < 1:
results_list.append(batch_outputs)
continue
batch_outputs = self._to_device(batch_outputs)
batch_outputs = self.model_forward(batch_outputs)
results = self.output_processor(batch_outputs)
if keep_inputs:
for i in range(len(batch_inputs)):
results[i].update({'inputs': batch_inputs[i]})
if isinstance(results, list):
results_list.extend(results)
else:
results_list.append(results)
prog_bar.update(len(batch_inputs))
# TODO: support append to file
if self.save_results:
self.dump(results_list, self.save_path)
return results_list