module.py 2.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
sys.path.insert(0, ".")

import time

T
Tingquan Gao 已提交
21
from paddlehub.utils.log import logger
22 23 24
from paddlehub.module.module import moduleinfo, serving
import cv2
import numpy as np
25
import paddle.nn as nn
26

T
Tingquan Gao 已提交
27 28
from tools.infer.predict import Predictor
from tools.infer.utils import b64_to_np, postprocess
29 30 31 32 33 34 35 36 37 38
from deploy.hubserving.clas.params import read_params


@moduleinfo(
    name="clas_system",
    version="1.0.0",
    summary="class system service",
    author="paddle-dev",
    author_email="paddle-dev@baidu.com",
    type="cv/class")
39 40
class ClasSystem(nn.Layer):
    def __init__(self, use_gpu=None, enable_mkldnn=None):
41 42 43 44 45 46
        """
        initialize with the necessary elements
        """
        cfg = read_params()
        if use_gpu is not None:
            cfg.use_gpu = use_gpu
T
Tingquan Gao 已提交
47 48
        if enable_mkldnn is not None:
            cfg.enable_mkldnn = enable_mkldnn
49 50 51 52 53 54 55 56 57 58 59 60 61 62 63
        cfg.hubserving = True
        cfg.enable_benchmark = False
        self.args = cfg
        if cfg.use_gpu:
            try:
                _places = os.environ["CUDA_VISIBLE_DEVICES"]
                int(_places[0])
                print("Use GPU, GPU Memery:{}".format(cfg.gpu_mem))
                print("CUDA_VISIBLE_DEVICES: ", _places)
            except:
                raise RuntimeError(
                    "Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id."
                )
        else:
            print("Use CPU")
T
Tingquan Gao 已提交
64
            print("Enable MKL-DNN") if enable_mkldnn else None
T
Tingquan Gao 已提交
65
        self.predictor = Predictor(self.args)
66

T
Tingquan Gao 已提交
67 68 69 70
    def predict(self, batch_input_data, top_k=1):
        assert isinstance(
            batch_input_data,
            np.ndarray), "The input data is inconsistent with expectations."
71

T
Tingquan Gao 已提交
72 73 74 75 76
        starttime = time.time()
        batch_outputs = self.predictor.predict(batch_input_data)
        elapse = time.time() - starttime
        batch_result_list = postprocess(batch_outputs, top_k)
        return {"prediction": batch_result_list, "elapse": elapse}
77 78

    @serving
T
Tingquan Gao 已提交
79
    def serving_method(self, images, revert_params, **kwargs):
80 81 82
        """
        Run as a service.
        """
T
Tingquan Gao 已提交
83 84
        input_data = b64_to_np(images, revert_params)
        results = self.predict(batch_input_data=input_data, **kwargs)
85
        return results