diff --git a/paddlehub/serving/app_single.py b/paddlehub/serving/app_single.py index a6975295c53e51b3535bd8255184bf0847635bb9..b1ddb1469e544f9e8a4a5e0851f0e9a0d277d8d0 100644 --- a/paddlehub/serving/app_single.py +++ b/paddlehub/serving/app_single.py @@ -14,6 +14,7 @@ from flask import Flask, request, render_template from paddlehub.serving.model_service.base_model_service import cv_module_info from paddlehub.serving.model_service.base_model_service import nlp_module_info +from paddlehub.serving.model_service.base_model_service import v2_module_info from paddlehub.common import utils import functools import time @@ -451,10 +452,8 @@ def create_app(init_flag=False, configs=None): @app_instance.route("/predict/", methods=["POST"]) def predict_modulev2(module_name): - if module_name in nlp_module_info.nlp_modules: - module_info = nlp_module_info.get_module_info(module_name) - elif module_name in cv_module_info.cv_modules: - module_info = cv_module_info.get_module_info(module_name) + if module_name in v2_module_info.modules: + module_info = v2_module_info.get_module_info(module_name) else: msg = "Module {} is not available.".format(module_name) return gen_result("-1", msg, "") @@ -476,6 +475,7 @@ def config_with_file(configs): cv_module_info.add_module(key, {key: value}) elif "NLP" == value["category"]: nlp_module_info.add_module(key, {key: value}) + v2_module_info.add_module(key, {key: value}) print(key, "==", value["version"]) diff --git a/paddlehub/serving/model_service/base_model_service.py b/paddlehub/serving/model_service/base_model_service.py index 803d4244e2c0005824203abb1e0b7f9b851cb120..586932288273e4af8750c0add798bed3817c6046 100644 --- a/paddlehub/serving/model_service/base_model_service.py +++ b/paddlehub/serving/model_service/base_model_service.py @@ -102,6 +102,19 @@ class NLPModuleInfo(BaseModuleInfo): self._modules.append(module_name) +class V2ModuleInfo(BaseModuleInfo): + def __init__(self): + super(V2ModuleInfo, self).__init__() + + @property + def modules(self): + return self._modules + + def add_module(self, module_name, module_info): + self._modules_info.update(module_info) + self._modules.append(module_name) + + class BaseModelService(object): def _initialize(self): pass @@ -121,3 +134,4 @@ class BaseModelService(object): cv_module_info = CVModuleInfo() nlp_module_info = NLPModuleInfo() +v2_module_info = V2ModuleInfo()