提交 658ec814 编写于 作者: C chenjian

reduce

上级 393c7d8f
...@@ -11,23 +11,29 @@ ...@@ -11,23 +11,29 @@
# 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 os
import argparse import argparse
import os
import paddle
import paddlehub as hub
from paddlehub.module.module import moduleinfo, runnable, serving
import numpy as np
import cv2 import cv2
import numpy as np
import paddle
import paddlehub as hub
from .enlighten_inference.pd_model.x2paddle_code import ONNXModel from .enlighten_inference.pd_model.x2paddle_code import ONNXModel
from .util import base64_to_cv2 from .util import base64_to_cv2
from paddlehub.module.module import moduleinfo
from paddlehub.module.module import runnable
from paddlehub.module.module import serving
@moduleinfo( @moduleinfo(name="enlightengan",
name="enlightengan", type="CV/enlighten", author="paddlepaddle", author_email="", summary="", version="1.0.0") type="CV/enlighten",
author="paddlepaddle",
author_email="",
summary="",
version="1.0.0")
class EnlightenGAN: class EnlightenGAN:
def __init__(self): def __init__(self):
self.pretrained_model = os.path.join(self.directory, "enlighten_inference/pd_model") self.pretrained_model = os.path.join(self.directory, "enlighten_inference/pd_model")
self.model = ONNXModel() self.model = ONNXModel()
...@@ -35,11 +41,11 @@ class EnlightenGAN: ...@@ -35,11 +41,11 @@ class EnlightenGAN:
self.model.set_dict(params, use_structured_name=True) self.model.set_dict(params, use_structured_name=True)
def enlightening(self, def enlightening(self,
images:list=None, images: list = None,
paths:list=None, paths: list = None,
output_dir:str='./enlightening_result/', output_dir: str = './enlightening_result/',
use_gpu:bool=False, use_gpu: bool = False,
visualization:bool=True): visualization: bool = True):
''' '''
enlighten images in the low-light scene. enlighten images in the low-light scene.
...@@ -95,11 +101,10 @@ class EnlightenGAN: ...@@ -95,11 +101,10 @@ class EnlightenGAN:
""" """
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(
...@@ -107,11 +112,10 @@ class EnlightenGAN: ...@@ -107,11 +112,10 @@ class EnlightenGAN:
self.add_module_config_arg() self.add_module_config_arg()
self.add_module_input_arg() self.add_module_input_arg()
self.args = self.parser.parse_args(argvs) self.args = self.parser.parse_args(argvs)
results = self.enlightening( results = self.enlightening(paths=[self.args.input_path],
paths=[self.args.input_path], output_dir=self.args.output_dir,
output_dir=self.args.output_dir, use_gpu=self.args.use_gpu,
use_gpu=self.args.use_gpu, visualization=self.args.visualization)
visualization=self.args.visualization)
return results return results
@serving @serving
...@@ -130,8 +134,10 @@ class EnlightenGAN: ...@@ -130,8 +134,10 @@ class EnlightenGAN:
""" """
self.arg_config_group.add_argument('--use_gpu', action='store_true', help="use GPU or not") self.arg_config_group.add_argument('--use_gpu', action='store_true', help="use GPU or not")
self.arg_config_group.add_argument( self.arg_config_group.add_argument('--output_dir',
'--output_dir', type=str, default='enlightening_result', help='output directory for saving result.') type=str,
default='enlightening_result',
help='output directory for saving result.')
self.arg_config_group.add_argument('--visualization', type=bool, default=False, help='save results or not.') self.arg_config_group.add_argument('--visualization', type=bool, default=False, help='save results or not.')
def add_module_input_arg(self): def add_module_input_arg(self):
......
import base64 import base64
import cv2 import cv2
import numpy as np import numpy as np
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册