|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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]) |
|
|
|
|
|
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 |
|
|