mirror of https://github.com/alibaba/EasyCV.git
332 lines
11 KiB
Python
332 lines
11 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import json
|
|
import os
|
|
import pickle
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from mmcv.parallel import collate, scatter_kwargs
|
|
from PIL import Image
|
|
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.mmlab_utils import (dynamic_adapt_for_mmlab,
|
|
remove_adapt_for_mmlab)
|
|
from easycv.utils.registry import build_from_cfg
|
|
|
|
|
|
class NumpyToPIL(object):
|
|
|
|
def __call__(self, results):
|
|
img = results['img']
|
|
results['img'] = Image.fromarray(np.uint8(img)).convert('RGB')
|
|
return results
|
|
|
|
|
|
class Predictor(object):
|
|
|
|
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)
|
|
|
|
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
|
|
|
|
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
|
|
|
|
|
|
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): Support 'cuda' or 'cpu', 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_IMAGE_MODE = 'BGR' # the image mode into the model
|
|
|
|
def __init__(self,
|
|
model_path,
|
|
config_file=None,
|
|
batch_size=1,
|
|
device=None,
|
|
save_results=False,
|
|
save_path=None,
|
|
pipelines=None,
|
|
*args,
|
|
**kwargs):
|
|
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
|
|
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"!')
|
|
|
|
self.model = self.prepare_model()
|
|
self.pipelines = pipelines
|
|
self.processor = self.build_processor()
|
|
self._load_op = None
|
|
|
|
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
|
|
|
|
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
|
|
|
|
def build_processor(self):
|
|
"""Build processor to process loaded input.
|
|
If you need custom preprocessing ops, you need to reimplement it.
|
|
"""
|
|
if self.pipelines is not None:
|
|
pipelines = self.pipelines
|
|
else:
|
|
pipelines = self.cfg.get('test_pipeline', [])
|
|
|
|
pipelines = [build_from_cfg(p, PIPELINES) for p in pipelines]
|
|
|
|
from easycv.datasets.shared.pipelines.transforms import Compose
|
|
processor = Compose(pipelines)
|
|
return processor
|
|
|
|
def _load_input(self, input):
|
|
"""Load image from file or numpy or PIL object.
|
|
Args:
|
|
input: File path or numpy or PIL object.
|
|
Returns:
|
|
{
|
|
'filename': filename,
|
|
'img': img,
|
|
'img_shape': img_shape,
|
|
'img_fields': ['img']
|
|
}
|
|
"""
|
|
if self._load_op is None:
|
|
load_cfg = dict(type='LoadImage', mode=self.INPUT_IMAGE_MODE)
|
|
self._load_op = build_from_cfg(load_cfg, PIPELINES)
|
|
|
|
if not isinstance(input, str):
|
|
if isinstance(input, np.ndarray):
|
|
# Only support RGB mode if input is np.ndarray.
|
|
input = cv2.cvtColor(input, cv2.COLOR_RGB2BGR)
|
|
sample = self._load_op({'img': input})
|
|
else:
|
|
sample = self._load_op({'filename': input})
|
|
|
|
return sample
|
|
|
|
def preprocess_single(self, input):
|
|
"""Preprocess single input sample.
|
|
If you need custom ops to load or process a single input sample, you need to reimplement it.
|
|
"""
|
|
input = self._load_input(input)
|
|
return self.processor(input)
|
|
|
|
def preprocess(self, inputs, *args, **kwargs):
|
|
"""Process all inputs list. And collate to batch and put to target device.
|
|
If you need custom ops to load or process a batch samples, you need to reimplement it.
|
|
"""
|
|
batch_outputs = []
|
|
for i in inputs:
|
|
batch_outputs.append(self.preprocess_single(i, *args, **kwargs))
|
|
|
|
batch_outputs = self._collate_fn(batch_outputs)
|
|
batch_outputs = self._to_device(batch_outputs)
|
|
|
|
return batch_outputs
|
|
|
|
def 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 postprocess(self, inputs, *args, **kwargs):
|
|
"""Process model batch outputs.
|
|
"""
|
|
outputs = []
|
|
out_i = {}
|
|
batch_size = 1
|
|
# get current batch size
|
|
for k, batch_v in inputs.items():
|
|
if batch_v is not None:
|
|
batch_size = len(batch_v)
|
|
break
|
|
|
|
for i in range(batch_size):
|
|
for k, batch_v in inputs.items():
|
|
if batch_v is not None:
|
|
out_i[k] = batch_v[i]
|
|
else:
|
|
out_i[k] = None
|
|
|
|
out_i = self.postprocess_single(out_i, *args, **kwargs)
|
|
outputs.append(out_i)
|
|
|
|
return outputs
|
|
|
|
def postprocess_single(self, inputs, *args, **kwargs):
|
|
"""Process outputs of single sample.
|
|
If you need add some processing ops, you need to reimplement it.
|
|
"""
|
|
return inputs
|
|
|
|
def _collate_fn(self, inputs):
|
|
"""Prepare the input just before the forward function.
|
|
Puts each data field into a tensor with outer dimension batch size
|
|
"""
|
|
return collate(inputs, samples_per_gpu=self.batch_size)
|
|
|
|
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]
|
|
|
|
@staticmethod
|
|
def dump(obj, save_path, mode='wb'):
|
|
with open(save_path, mode) as f:
|
|
f.write(pickle.dumps(obj))
|
|
|
|
def __call__(self, inputs, keep_inputs=False):
|
|
# TODO: fault tolerance
|
|
|
|
if isinstance(inputs, str):
|
|
inputs = [inputs]
|
|
|
|
results_list = []
|
|
for i in range(0, len(inputs), self.batch_size):
|
|
batch = inputs[i:min(len(inputs), i + self.batch_size)]
|
|
batch_outputs = self.preprocess(batch)
|
|
batch_outputs = self.forward(batch_outputs)
|
|
results = self.postprocess(batch_outputs)
|
|
# assert len(results) == len(
|
|
# batch), f'Mismatch size {len(results)} != {len(batch)}'
|
|
if keep_inputs:
|
|
for i in range(len(batch)):
|
|
results[i].update({'inputs': batch[i]})
|
|
# if dump, the outputs will not added to the return value to prevent taking up too much memory
|
|
if self.save_results:
|
|
self.dump(results, self.save_path, mode='ab+')
|
|
else:
|
|
if isinstance(results, list):
|
|
results_list.extend(results)
|
|
else:
|
|
results_list.append(results)
|
|
|
|
return results_list
|