Skip to content

P
PaddleHub

PaddlePaddle / PaddleHub
大约 1 年 前同步成功

通知 282
Star 12117
Fork 2091
P
PaddleHub
未验证 提交 e301f567 编写于 作者: jm_12138's avatar jm_12138 提交者： GitHub
浏览文件
操作

update humanseg_lite (#2181)

上级 3b20f676
隐藏空白更改
内联 并排
Showing with 136 addition and 160 deletion
modules/image/semantic_segmentation/humanseg_lite/README.md
浏览文件 @ e301f567
# humanseg_lite # humanseg_lite
|模型名称|humanseg_lite| |模型名称|humanseg_lite|
| :--- | :---: | | :--- | :---: |
|类别|图像-图像分割| |类别|图像-图像分割|
|网络|shufflenet| |网络|shufflenet|
|数据集|百度自建数据集| |数据集|百度自建数据集|
...@@ -13,7 +13,7 @@ ...@@ -13,7 +13,7 @@
## 一、模型基本信息 ## 一、模型基本信息
- ### 应用效果展示 - ### 应用效果展示
- 样例结果示例： - 样例结果示例：
<p align="center"> <p align="center">
<img src="https://user-images.githubusercontent.com/35907364/130913092-312a5f37-842e-4fd0-8db4-5f853fd8419f.jpg" width = "337" height = "505" hspace='10'/> <img src="https://user-images.githubusercontent.com/35907364/130916087-7d537ad9-bbc8-4bce-9382-8eb132b35532.png" width = "337" height = "505" hspace='10'/> <img src="https://user-images.githubusercontent.com/35907364/130913092-312a5f37-842e-4fd0-8db4-5f853fd8419f.jpg" width = "337" height = "505" hspace='10'/> <img src="https://user-images.githubusercontent.com/35907364/130916087-7d537ad9-bbc8-4bce-9382-8eb132b35532.png" width = "337" height = "505" hspace='10'/>
...@@ -37,7 +37,7 @@ ...@@ -37,7 +37,7 @@
- ```shell - ```shell
$ hub install humanseg_lite $ hub install humanseg_lite
``` ```
- 如您安装时遇到问题，可参考：[零基础windows安装](../../../../docs/docs_ch/get_start/windows_quickstart.md) - 如您安装时遇到问题，可参考：[零基础windows安装](../../../../docs/docs_ch/get_start/windows_quickstart.md)
| [零基础Linux安装](../../../../docs/docs_ch/get_start/linux_quickstart.md) | [零基础MacOS安装](../../../../docs/docs_ch/get_start/mac_quickstart.md) | [零基础Linux安装](../../../../docs/docs_ch/get_start/linux_quickstart.md) | [零基础MacOS安装](../../../../docs/docs_ch/get_start/mac_quickstart.md)
...@@ -226,6 +226,9 @@ ...@@ -226,6 +226,9 @@
cv2.imwrite("segment_human_lite.png", rgba) cv2.imwrite("segment_human_lite.png", rgba)
``` ```
- ### Gradio APP 支持
从 PaddleHub 2.3.1 开始支持使用链接 http://127.0.0.1:8866/gradio/humanseg_lite 在浏览器中访问 humanseg_lite 的 Gradio APP。
## 五、更新历史 ## 五、更新历史
...@@ -234,7 +237,7 @@ ...@@ -234,7 +237,7 @@
初始发布 初始发布
* 1.1.0 * 1.1.0
新增视频人像分割接口 新增视频人像分割接口
新增视频流人像分割接口 新增视频流人像分割接口
...@@ -247,6 +250,10 @@ ...@@ -247,6 +250,10 @@
移除 Fluid API 移除 Fluid API
* 1.3.0
添加 Gradio APP 支持
```shell ```shell
$ hub install humanseg_lite == 1.2.0 $ hub install humanseg_lite == 1.3.0
``` ```
modules/image/semantic_segmentation/humanseg_lite/README_en.md
浏览文件 @ e301f567
# humanseg_lite # humanseg_lite
|Module Name |humanseg_lite| |Module Name |humanseg_lite|
| :--- | :---: | | :--- | :---: |
|Category |Image segmentation| |Category |Image segmentation|
|Network|shufflenet| |Network|shufflenet|
|Dataset|Baidu self-built dataset| |Dataset|Baidu self-built dataset|
...@@ -10,10 +10,10 @@ ...@@ -10,10 +10,10 @@
|Data indicators|-| |Data indicators|-|
|Latest update date|2021-02-26| |Latest update date|2021-02-26|
## I. Basic Information ## I. Basic Information
- ### Application Effect Display - ### Application Effect Display
- Sample results: - Sample results:
<p align="center"> <p align="center">
<img src="https://user-images.githubusercontent.com/35907364/130913092-312a5f37-842e-4fd0-8db4-5f853fd8419f.jpg" width = "337" height = "505" hspace='10'/> <img src="https://user-images.githubusercontent.com/35907364/130916087-7d537ad9-bbc8-4bce-9382-8eb132b35532.png" width = "337" height = "505" hspace='10'/> <img src="https://user-images.githubusercontent.com/35907364/130913092-312a5f37-842e-4fd0-8db4-5f853fd8419f.jpg" width = "337" height = "505" hspace='10'/> <img src="https://user-images.githubusercontent.com/35907364/130916087-7d537ad9-bbc8-4bce-9382-8eb132b35532.png" width = "337" height = "505" hspace='10'/>
...@@ -39,7 +39,7 @@ ...@@ -39,7 +39,7 @@
- ```shell - ```shell
$ hub install humanseg_lite $ hub install humanseg_lite
``` ```
- In case of any problems during installation, please refer to:[Windows_Quickstart](../../../../docs/docs_en/get_start/windows_quickstart.md) - In case of any problems during installation, please refer to:[Windows_Quickstart](../../../../docs/docs_en/get_start/windows_quickstart.md)
| [Linux_Quickstart](../../../../docs/docs_en/get_start/linux_quickstart.md) | [Mac_Quickstart](../../../../docs/docs_en/get_start/mac_quickstart.md) | [Linux_Quickstart](../../../../docs/docs_en/get_start/linux_quickstart.md) | [Mac_Quickstart](../../../../docs/docs_en/get_start/mac_quickstart.md)
...@@ -49,7 +49,7 @@ ...@@ -49,7 +49,7 @@
- ``` - ```
hub run humanseg_lite --input_path "/PATH/TO/IMAGE" hub run humanseg_lite --input_path "/PATH/TO/IMAGE"
``` ```
- If you want to call the Hub module through the command line, please refer to: [PaddleHub Command Line Instruction](../../../../docs/docs_en/tutorial/cmd_usage.rst) - If you want to call the Hub module through the command line, please refer to: [PaddleHub Command Line Instruction](../../../../docs/docs_en/tutorial/cmd_usage.rst)
...@@ -122,9 +122,9 @@ ...@@ -122,9 +122,9 @@
- **Return** - **Return**
* res (list\[dict\]): The list of recognition results, where each element is dict and each field is: * res (list\[dict\]): The list of recognition results, where each element is dict and each field is:
* save\_path (str, optional): Save path of the result. * save\_path (str, optional): Save path of the result.
* data (numpy.ndarray): The result of portrait segmentation. * data (numpy.ndarray): The result of portrait segmentation.
- ```python - ```python
def video_stream_segment(self, def video_stream_segment(self,
...@@ -230,6 +230,8 @@ ...@@ -230,6 +230,8 @@
cv2.imwrite("segment_human_lite.png", rgba) cv2.imwrite("segment_human_lite.png", rgba)
``` ```
- ### Gradio APP support
Starting with PaddleHub 2.3.1, the Gradio APP for humanseg_lite is supported to be accessed in the browser using the link http://127.0.0.1:8866/gradio/humanseg_lite.
## V. Release Note ## V. Release Note
...@@ -238,7 +240,7 @@ ...@@ -238,7 +240,7 @@
First release First release
- 1.1.0 - 1.1.0
Added video portrait segmentation interface Added video portrait segmentation interface
Added video stream portrait segmentation interface Added video stream portrait segmentation interface
...@@ -251,6 +253,10 @@ ...@@ -251,6 +253,10 @@
Remove Fluid API Remove Fluid API
* 1.3.0
Add Gradio APP support.
```shell ```shell
$ hub install humanseg_lite == 1.2.0 $ hub install humanseg_lite == 1.3.0
``` ```
modules/image/semantic_segmentation/humanseg_lite/module.py
浏览文件 @ e301f567
...@@ -12,32 +12,39 @@ ...@@ -12,32 +12,39 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import argparse
import ast import ast
import os import os
import os.path as osp import os.path as osp
import argparse
import cv2 import cv2
import numpy as np import numpy as np
import paddle
import paddle.jit import paddle.jit
import paddle.static import paddle.static
from paddle.inference import Config, create_predictor from paddle.inference import Config
from paddlehub.module.module import moduleinfo, runnable, serving from paddle.inference import create_predictor
from .processor import postprocess, base64_to_cv2, cv2_to_base64, check_dir from .data_feed import preprocess_v
from .data_feed import reader, preprocess_v from .data_feed import reader
from .optimal import postprocess_v, threshold_mask from .optimal import postprocess_v
from .optimal import threshold_mask
from .processor import base64_to_cv2
@moduleinfo( from .processor import check_dir
name="humanseg_lite", from .processor import cv2_to_base64
type="CV/semantic_segmentation", from .processor import postprocess
author="paddlepaddle", from paddlehub.module.module import moduleinfo
author_email="", from paddlehub.module.module import runnable
summary="humanseg_lite is a semantic segmentation model.", from paddlehub.module.module import serving
version="1.2.0")
@moduleinfo(name="humanseg_lite",
type="CV/semantic_segmentation",
author="paddlepaddle",
author_email="",
summary="humanseg_lite is a semantic segmentation model.",
version="1.3.0")
class ShufflenetHumanSeg: class ShufflenetHumanSeg:
def __init__(self): def __init__(self):
self.default_pretrained_model_path = os.path.join(self.directory, "humanseg_lite_inference", "model") self.default_pretrained_model_path = os.path.join(self.directory, "humanseg_lite_inference", "model")
self._set_config() self._set_config()
...@@ -46,8 +53,8 @@ class ShufflenetHumanSeg: ...@@ -46,8 +53,8 @@ class ShufflenetHumanSeg:
""" """
predictor config setting predictor config setting
""" """
model = self.default_pretrained_model_path+'.pdmodel' model = self.default_pretrained_model_path + '.pdmodel'
params = self.default_pretrained_model_path+'.pdiparams' params = self.default_pretrained_model_path + '.pdiparams'
cpu_config = Config(model, params) cpu_config = Config(model, params)
cpu_config.disable_glog_info() cpu_config.disable_glog_info()
cpu_config.disable_gpu() cpu_config.disable_gpu()
...@@ -64,7 +71,7 @@ class ShufflenetHumanSeg: ...@@ -64,7 +71,7 @@ class ShufflenetHumanSeg:
gpu_config = Config(model, params) gpu_config = Config(model, params)
gpu_config.disable_glog_info() gpu_config.disable_glog_info()
gpu_config.enable_use_gpu(memory_pool_init_size_mb=1000, device_id=0) gpu_config.enable_use_gpu(memory_pool_init_size_mb=1000, device_id=0)
if paddle.get_cudnn_version() == 8004: if paddle.get_cudnn_version() == 8004:
gpu_config.delete_pass('conv_elementwise_add_act_fuse_pass') gpu_config.delete_pass('conv_elementwise_add_act_fuse_pass')
gpu_config.delete_pass('conv_elementwise_add2_act_fuse_pass') gpu_config.delete_pass('conv_elementwise_add2_act_fuse_pass')
...@@ -134,13 +141,12 @@ class ShufflenetHumanSeg: ...@@ -134,13 +141,12 @@ class ShufflenetHumanSeg:
output = np.expand_dims(output[:, 1, :, :], axis=1) output = np.expand_dims(output[:, 1, :, :], axis=1)
# postprocess one by one # postprocess one by one
for i in range(len(batch_data)): for i in range(len(batch_data)):
out = postprocess( out = postprocess(data_out=output[i],
data_out=output[i], org_im=batch_data[i]['org_im'],
org_im=batch_data[i]['org_im'], org_im_shape=batch_data[i]['org_im_shape'],
org_im_shape=batch_data[i]['org_im_shape'], org_im_path=batch_data[i]['org_im_path'],
org_im_path=batch_data[i]['org_im_path'], output_dir=output_dir,
output_dir=output_dir, visualization=visualization)
visualization=visualization)
res.append(out) res.append(out)
return res return res
...@@ -327,11 +333,10 @@ class ShufflenetHumanSeg: ...@@ -327,11 +333,10 @@ class ShufflenetHumanSeg:
""" """
Run as a command. Run as a command.
""" """
self.parser = argparse.ArgumentParser( self.parser = argparse.ArgumentParser(description="Run the {} module.".format(self.name),
description="Run the {} module.".format(self.name), prog='hub run {}'.format(self.name),
prog='hub run {}'.format(self.name), usage='%(prog)s',
usage='%(prog)s', add_help=True)
add_help=True)
self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required") self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required")
self.arg_config_group = self.parser.add_argument_group( self.arg_config_group = self.parser.add_argument_group(
...@@ -339,12 +344,11 @@ class ShufflenetHumanSeg: ...@@ -339,12 +344,11 @@ class ShufflenetHumanSeg:
self.add_module_config_arg() self.add_module_config_arg()
self.add_module_input_arg() self.add_module_input_arg()
args = self.parser.parse_args(argvs) args = self.parser.parse_args(argvs)
results = self.segment( results = self.segment(paths=[args.input_path],
paths=[args.input_path], batch_size=args.batch_size,
batch_size=args.batch_size, use_gpu=args.use_gpu,
use_gpu=args.use_gpu, output_dir=args.output_dir,
output_dir=args.output_dir, visualization=args.visualization)
visualization=args.visualization)
if args.save_dir is not None: if args.save_dir is not None:
check_dir(args.save_dir) check_dir(args.save_dir)
self.save_inference_model(args.save_dir) self.save_inference_model(args.save_dir)
...@@ -355,14 +359,22 @@ class ShufflenetHumanSeg: ...@@ -355,14 +359,22 @@ class ShufflenetHumanSeg:
""" """
Add the command config options. Add the command config options.
""" """
self.arg_config_group.add_argument( self.arg_config_group.add_argument('--use_gpu',
'--use_gpu', type=ast.literal_eval, default=False, help="whether use GPU or not") type=ast.literal_eval,
self.arg_config_group.add_argument( default=False,
'--output_dir', type=str, default='humanseg_lite_output', help="The directory to save output images.") help="whether use GPU or not")
self.arg_config_group.add_argument( self.arg_config_group.add_argument('--output_dir',
'--save_dir', type=str, default='humanseg_lite_model', help="The directory to save model.") type=str,
self.arg_config_group.add_argument( default='humanseg_lite_output',
'--visualization', type=ast.literal_eval, default=False, help="whether to save output as images.") help="The directory to save output images.")
self.arg_config_group.add_argument('--save_dir',
type=str,
default='humanseg_lite_model',
help="The directory to save model.")
self.arg_config_group.add_argument('--visualization',
type=ast.literal_eval,
default=False,
help="whether to save output as images.")
self.arg_config_group.add_argument('--batch_size', type=ast.literal_eval, default=1, help="batch size.") self.arg_config_group.add_argument('--batch_size', type=ast.literal_eval, default=1, help="batch size.")
def add_module_input_arg(self): def add_module_input_arg(self):
...@@ -371,33 +383,20 @@ class ShufflenetHumanSeg: ...@@ -371,33 +383,20 @@ class ShufflenetHumanSeg:
""" """
self.arg_input_group.add_argument('--input_path', type=str, help="path to image.") self.arg_input_group.add_argument('--input_path', type=str, help="path to image.")
def create_gradio_app(self):
if __name__ == "__main__": import gradio as gr
m = ShufflenetHumanSeg() import tempfile
#shuffle.video_segment() import os
img = cv2.imread('photo.jpg') from PIL import Image
# res = m.segment(images=[img], visualization=True)
# print(res[0]['data']) def inference(image, use_gpu=False):
# m.video_segment('') with tempfile.TemporaryDirectory() as temp_dir:
cap_video = cv2.VideoCapture('video_test.mp4') self.segment(paths=[image], use_gpu=use_gpu, visualization=True, output_dir=temp_dir)
fps = cap_video.get(cv2.CAP_PROP_FPS) return Image.open(os.path.join(temp_dir, os.listdir(temp_dir)[0]))
save_path = 'result_frame.avi'
width = int(cap_video.get(cv2.CAP_PROP_FRAME_WIDTH)) interface = gr.Interface(
height = int(cap_video.get(cv2.CAP_PROP_FRAME_HEIGHT)) inference,
cap_out = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (width, height)) [gr.inputs.Image(type="filepath"), gr.Checkbox(label='use_gpu')],
prev_gray = None gr.outputs.Image(type="ndarray"),
prev_cfd = None title='humanseg_lite')
while cap_video.isOpened(): return interface
ret, frame_org = cap_video.read()
if ret:
[img_matting, prev_gray, prev_cfd] = m.video_stream_segment(
frame_org=frame_org, frame_id=cap_video.get(1), prev_gray=prev_gray, prev_cfd=prev_cfd)
img_matting = np.repeat(img_matting[:, :, np.newaxis], 3, axis=2)
bg_im = np.ones_like(img_matting) * 255
comb = (img_matting * frame_org + (1 - img_matting) * bg_im).astype(np.uint8)
cap_out.write(comb)
else:
break
cap_video.release()
cap_out.release()
modules/image/semantic_segmentation/humanseg_lite/optimal.py
浏览文件 @ e301f567
...@@ -32,8 +32,8 @@ def human_seg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow): ...@@ -32,8 +32,8 @@ def human_seg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow):
# 超出边界不跟踪 # 超出边界不跟踪
not_track = (cur_x < 0) + (cur_x >= w) + (cur_y < 0) + (cur_y >= h) not_track = (cur_x < 0) + (cur_x >= w) + (cur_y < 0) + (cur_y >= h)
flow_bw[~not_track] = flow_bw[cur_y[~not_track], cur_x[~not_track]] flow_bw[~not_track] = flow_bw[cur_y[~not_track], cur_x[~not_track]]
not_track += ( not_track += (np.square(flow_fw[:, :, 0] + flow_bw[:, :, 0]) +
np.square(flow_fw[:, :, 0] + flow_bw[:, :, 0]) + np.square(flow_fw[:, :, 1] + flow_bw[:, :, 1])) >= check_thres np.square(flow_fw[:, :, 1] + flow_bw[:, :, 1])) >= check_thres
track_cfd[cur_y[~not_track], cur_x[~not_track]] = prev_cfd[~not_track] track_cfd[cur_y[~not_track], cur_x[~not_track]] = prev_cfd[~not_track]
is_track[cur_y[~not_track], cur_x[~not_track]] = 1 is_track[cur_y[~not_track], cur_x[~not_track]] = 1
......
modules/image/semantic_segmentation/humanseg_lite/processor.py
浏览文件 @ e301f567
# -*- coding:utf-8 -*- # -*- coding:utf-8 -*-
import base64
import os import os
import time import time
import base64
import cv2 import cv2
import numpy as np import numpy as np
......
modules/image/semantic_segmentation/humanseg_lite/test.py
浏览文件 @ e301f567
...@@ -3,15 +3,16 @@ import shutil ...@@ -3,15 +3,16 @@ import shutil
import unittest import unittest
import cv2 import cv2
import requests
import numpy as np import numpy as np
import paddlehub as hub import requests
import paddlehub as hub
os.environ['CUDA_VISIBLE_DEVICES'] = '0' os.environ['CUDA_VISIBLE_DEVICES'] = '0'
class TestHubModule(unittest.TestCase): class TestHubModule(unittest.TestCase):
@classmethod @classmethod
def setUpClass(cls) -> None: def setUpClass(cls) -> None:
img_url = 'https://unsplash.com/photos/pg_WCHWSdT8/download?ixid=MnwxMjA3fDB8MXxhbGx8fHx8fHx8fHwxNjYyNDM2ODI4&force=true&w=640' img_url = 'https://unsplash.com/photos/pg_WCHWSdT8/download?ixid=MnwxMjA3fDB8MXxhbGx8fHx8fHx8fHwxNjYyNDM2ODI4&force=true&w=640'
...@@ -23,8 +24,7 @@ class TestHubModule(unittest.TestCase): ...@@ -23,8 +24,7 @@ class TestHubModule(unittest.TestCase):
f.write(response.content) f.write(response.content)
fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G') fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
img = cv2.imread('tests/test.jpg') img = cv2.imread('tests/test.jpg')
video = cv2.VideoWriter('tests/test.avi', fourcc, video = cv2.VideoWriter('tests/test.avi', fourcc, 20.0, tuple(img.shape[:2]))
20.0, tuple(img.shape[:2]))
for i in range(40): for i in range(40):
video.write(img) video.write(img)
video.release() video.release()
...@@ -38,101 +38,65 @@ class TestHubModule(unittest.TestCase): ...@@ -38,101 +38,65 @@ class TestHubModule(unittest.TestCase):
shutil.rmtree('humanseg_lite_video_result') shutil.rmtree('humanseg_lite_video_result')
def test_segment1(self): def test_segment1(self):
results = self.module.segment( results = self.module.segment(paths=['tests/test.jpg'], use_gpu=False, visualization=False)
paths=['tests/test.jpg'],
use_gpu=False,
visualization=False
)
self.assertIsInstance(results[0]['data'], np.ndarray) self.assertIsInstance(results[0]['data'], np.ndarray)
def test_segment2(self): def test_segment2(self):
results = self.module.segment( results = self.module.segment(images=[cv2.imread('tests/test.jpg')], use_gpu=False, visualization=False)
images=[cv2.imread('tests/test.jpg')],
use_gpu=False,
visualization=False
)
self.assertIsInstance(results[0]['data'], np.ndarray) self.assertIsInstance(results[0]['data'], np.ndarray)
def test_segment3(self): def test_segment3(self):
results = self.module.segment( results = self.module.segment(images=[cv2.imread('tests/test.jpg')], use_gpu=False, visualization=True)
images=[cv2.imread('tests/test.jpg')],
use_gpu=False,
visualization=True
)
self.assertIsInstance(results[0]['data'], np.ndarray) self.assertIsInstance(results[0]['data'], np.ndarray)
def test_segment4(self): def test_segment4(self):
results = self.module.segment( results = self.module.segment(images=[cv2.imread('tests/test.jpg')], use_gpu=True, visualization=False)
images=[cv2.imread('tests/test.jpg')],
use_gpu=True,
visualization=False
)
self.assertIsInstance(results[0]['data'], np.ndarray) self.assertIsInstance(results[0]['data'], np.ndarray)
def test_segment5(self): def test_segment5(self):
self.assertRaises( self.assertRaises(AssertionError, self.module.segment, paths=['no.jpg'])
AssertionError,
self.module.segment,
paths=['no.jpg']
)
def test_segment6(self): def test_segment6(self):
self.assertRaises( self.assertRaises(AttributeError, self.module.segment, images=['test.jpg'])
AttributeError,
self.module.segment,
images=['test.jpg']
)
def test_video_stream_segment1(self): def test_video_stream_segment1(self):
img_matting, cur_gray, optflow_map = self.module.video_stream_segment( img_matting, cur_gray, optflow_map = self.module.video_stream_segment(frame_org=cv2.imread('tests/test.jpg'),
frame_org=cv2.imread('tests/test.jpg'), frame_id=1,
frame_id=1, prev_gray=None,
prev_gray=None, prev_cfd=None,
prev_cfd=None, use_gpu=False)
use_gpu=False
)
self.assertIsInstance(img_matting, np.ndarray) self.assertIsInstance(img_matting, np.ndarray)
self.assertIsInstance(cur_gray, np.ndarray) self.assertIsInstance(cur_gray, np.ndarray)
self.assertIsInstance(optflow_map, np.ndarray) self.assertIsInstance(optflow_map, np.ndarray)
img_matting, cur_gray, optflow_map = self.module.video_stream_segment( img_matting, cur_gray, optflow_map = self.module.video_stream_segment(frame_org=cv2.imread('tests/test.jpg'),
frame_org=cv2.imread('tests/test.jpg'), frame_id=2,
frame_id=2, prev_gray=cur_gray,
prev_gray=cur_gray, prev_cfd=optflow_map,
prev_cfd=optflow_map, use_gpu=False)
use_gpu=False
)
self.assertIsInstance(img_matting, np.ndarray) self.assertIsInstance(img_matting, np.ndarray)
self.assertIsInstance(cur_gray, np.ndarray) self.assertIsInstance(cur_gray, np.ndarray)
self.assertIsInstance(optflow_map, np.ndarray) self.assertIsInstance(optflow_map, np.ndarray)
def test_video_stream_segment2(self): def test_video_stream_segment2(self):
img_matting, cur_gray, optflow_map = self.module.video_stream_segment( img_matting, cur_gray, optflow_map = self.module.video_stream_segment(frame_org=cv2.imread('tests/test.jpg'),
frame_org=cv2.imread('tests/test.jpg'), frame_id=1,
frame_id=1, prev_gray=None,
prev_gray=None, prev_cfd=None,
prev_cfd=None, use_gpu=True)
use_gpu=True
)
self.assertIsInstance(img_matting, np.ndarray) self.assertIsInstance(img_matting, np.ndarray)
self.assertIsInstance(cur_gray, np.ndarray) self.assertIsInstance(cur_gray, np.ndarray)
self.assertIsInstance(optflow_map, np.ndarray) self.assertIsInstance(optflow_map, np.ndarray)
img_matting, cur_gray, optflow_map = self.module.video_stream_segment( img_matting, cur_gray, optflow_map = self.module.video_stream_segment(frame_org=cv2.imread('tests/test.jpg'),
frame_org=cv2.imread('tests/test.jpg'), frame_id=2,
frame_id=2, prev_gray=cur_gray,
prev_gray=cur_gray, prev_cfd=optflow_map,
prev_cfd=optflow_map, use_gpu=True)
use_gpu=True
)
self.assertIsInstance(img_matting, np.ndarray) self.assertIsInstance(img_matting, np.ndarray)
self.assertIsInstance(cur_gray, np.ndarray) self.assertIsInstance(cur_gray, np.ndarray)
self.assertIsInstance(optflow_map, np.ndarray) self.assertIsInstance(optflow_map, np.ndarray)
def test_video_segment1(self): def test_video_segment1(self):
self.module.video_segment( self.module.video_segment(video_path="tests/test.avi", use_gpu=False, save_dir='humanseg_lite_video_result')
video_path="tests/test.avi",
use_gpu=False,
save_dir='humanseg_lite_video_result'
)
def test_save_inference_model(self): def test_save_inference_model(self):
self.module.save_inference_model('./inference/model') self.module.save_inference_model('./inference/model')
......
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或