提交 dae3b319 编写于 作者: C chenjian

modify according to review

上级 b645cf17
...@@ -31,6 +31,9 @@ from .util import base64_to_cv2 ...@@ -31,6 +31,9 @@ from .util import base64_to_cv2
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()
params = paddle.load(os.path.join(self.pretrained_model, 'model.pdparams'))
self.model.set_dict(params, use_structured_name=True)
def enlightening(self, def enlightening(self,
images=None, images=None,
...@@ -54,9 +57,6 @@ class EnlightenGAN: ...@@ -54,9 +57,6 @@ class EnlightenGAN:
if images == None and paths == None: if images == None and paths == None:
print('No image provided. Please input an image or a image path.') print('No image provided. Please input an image or a image path.')
return return
self.model = ONNXModel()
params = paddle.load(os.path.join(self.pretrained_model, 'model.pdparams'))
self.model.set_dict(params, use_structured_name=True)
self.model.eval() self.model.eval()
if images != None: if images != None:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册