radames's picture
first
a0bcaae
raw
history blame
6.33 kB
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import codecs
import os
import time
import yaml
import numpy as np
import cv2
import paddle
import paddleseg.transforms as T
from paddle.inference import create_predictor, PrecisionType
from paddle.inference import Config as PredictConfig
from paddleseg.core.infer import reverse_transform
from paddleseg.cvlibs import manager
from paddleseg.utils import TimeAverager
from ..scripts.optic_flow_process import optic_flow_process
class DeployConfig:
def __init__(self, path):
with codecs.open(path, 'r', 'utf-8') as file:
self.dic = yaml.load(file, Loader=yaml.FullLoader)
self._transforms = self._load_transforms(self.dic['Deploy'][
'transforms'])
self._dir = os.path.dirname(path)
@property
def transforms(self):
return self._transforms
@property
def model(self):
return os.path.join(self._dir, self.dic['Deploy']['model'])
@property
def params(self):
return os.path.join(self._dir, self.dic['Deploy']['params'])
def _load_transforms(self, t_list):
com = manager.TRANSFORMS
transforms = []
for t in t_list:
ctype = t.pop('type')
transforms.append(com[ctype](**t))
return transforms
class Predictor:
def __init__(self, args):
self.cfg = DeployConfig(args.cfg)
self.args = args
self.compose = T.Compose(self.cfg.transforms)
resize_h, resize_w = args.input_shape
self.disflow = cv2.DISOpticalFlow_create(
cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST)
self.prev_gray = np.zeros((resize_h, resize_w), np.uint8)
self.prev_cfd = np.zeros((resize_h, resize_w), np.float32)
self.is_init = True
pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
pred_cfg.disable_glog_info()
if self.args.use_gpu:
pred_cfg.enable_use_gpu(100, 0)
self.predictor = create_predictor(pred_cfg)
if self.args.test_speed:
self.cost_averager = TimeAverager()
def preprocess(self, img):
ori_shapes = []
processed_imgs = []
processed_img = self.compose(img)[0]
processed_imgs.append(processed_img)
ori_shapes.append(img.shape)
return processed_imgs, ori_shapes
def run(self, img, bg):
input_names = self.predictor.get_input_names()
input_handle = self.predictor.get_input_handle(input_names[0])
processed_imgs, ori_shapes = self.preprocess(img)
data = np.array(processed_imgs)
input_handle.reshape(data.shape)
input_handle.copy_from_cpu(data)
if self.args.test_speed:
start = time.time()
self.predictor.run()
if self.args.test_speed:
self.cost_averager.record(time.time() - start)
output_names = self.predictor.get_output_names()
output_handle = self.predictor.get_output_handle(output_names[0])
output = output_handle.copy_to_cpu()
return self.postprocess(output, img, ori_shapes[0], bg)
def postprocess(self, pred, img, ori_shape, bg):
if not os.path.exists(self.args.save_dir):
os.makedirs(self.args.save_dir)
resize_w = pred.shape[-1]
resize_h = pred.shape[-2]
if self.args.soft_predict:
if self.args.use_optic_flow:
score_map = pred[:, 1, :, :].squeeze(0)
score_map = 255 * score_map
cur_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
cur_gray = cv2.resize(cur_gray, (resize_w, resize_h))
optflow_map = optic_flow_process(cur_gray, score_map, self.prev_gray, self.prev_cfd, \
self.disflow, self.is_init)
self.prev_gray = cur_gray.copy()
self.prev_cfd = optflow_map.copy()
self.is_init = False
score_map = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2)
score_map = np.transpose(score_map, [2, 0, 1])[np.newaxis, ...]
score_map = reverse_transform(
paddle.to_tensor(score_map),
ori_shape,
self.cfg.transforms,
mode='bilinear')
alpha = np.transpose(score_map.numpy().squeeze(0),
[1, 2, 0]) / 255
else:
score_map = pred[:, 1, :, :]
score_map = score_map[np.newaxis, ...]
score_map = reverse_transform(
paddle.to_tensor(score_map),
ori_shape,
self.cfg.transforms,
mode='bilinear')
alpha = np.transpose(score_map.numpy().squeeze(0), [1, 2, 0])
else:
if pred.ndim == 3:
pred = pred[:, np.newaxis, ...]
result = reverse_transform(
paddle.to_tensor(
pred, dtype='float32'),
ori_shape,
self.cfg.transforms,
mode='bilinear')
result = np.array(result)
if self.args.add_argmax:
result = np.argmax(result, axis=1)
else:
result = result.squeeze(1)
alpha = np.transpose(result, [1, 2, 0])
# background replace
h, w, _ = img.shape
if bg is None:
bg = np.ones_like(img)*255
else:
bg = cv2.resize(bg, (w, h))
if bg.ndim == 2:
bg = bg[..., np.newaxis]
comb = (alpha * img + (1 - alpha) * bg).astype(np.uint8)
return comb, alpha, bg, img