未验证 提交 5acac161 编写于 作者: H haoyuying 提交者: GitHub

Fix the problem that the colorization model reads data incorrectly on Paddle 2.1.

上级 778db1ae
......@@ -31,7 +31,7 @@ import paddlehub.vision.transforms as T
import paddlehub.vision.functional as Func
from paddlehub.vision import utils
from paddlehub.module.module import serving, RunModule, runnable
from paddlehub.utils.utils import base64_to_cv2, cv2_to_base64
from paddlehub.utils.utils import base64_to_cv2, cv2_to_base64, Version
class ImageServing(object):
......@@ -203,7 +203,11 @@ class ImageColorizeModule(RunModule, ImageServing):
Returns:
results(dict) : The model outputs, such as metrics.
'''
img = self.preprocess(batch[0])
if Version(paddle.__version__) >= '2.1':
img = self.preprocess(batch)
else:
img = self.preprocess(batch[0])
out_class, out_reg = self(img['A'], img['hint_B'], img['mask_B'])
# loss
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册