server_develop/code/src/object_detector_handler.py

# -*- coding: utf-8 -*-
# @Time : 2022/2/28 11:29
# @Author : MaochengHu
# @Email : wojiaohumaocheng@gmail.com
# @File : object_detector_handler.py
# @Project : server_develop
import os
import sys
import logging

def import_env(names):
    for name in names:
        sys.path.append(os.path.join(os.getcwd(), name))


names = ["src", "test", "config_files", "src/model_structure", "src/server_utils",
         "test", "test/test_src", "src/model_structure/yolov5_structure",
         "src/model_structure/yolov5_structure/yolov5-master",
         "src/model_structure/yolov5_structure/yolov5-master/utils/"]
import_env(names)



from server_utils.method_parser import parse
from server_utils.recorder import Timer, logger
from server_utils.image_convert import base642tensor
from basic_model_handler import BasicModelHandler
from model_structure.yolov5_structure.yolov5 import YOLOV5


class ObjectDetectionHandler(BasicModelHandler):
    def __init__(self):
        super(ObjectDetectionHandler, self).__init__()
        self.yolov5_model = YOLOV5().create()

    @Timer.time_recorder
    def preprocess(self, data, cuda=True):
        body_data = data.get("body", None)
        assert body_data is not None, "get body(request info) is None"
        input_data = body_data.get("data")
        assert body_data is not None, "get data(request info) is None"
        output_methods = body_data.get("output_methods", [])
        input_source = base642tensor(input_data)
        logger.debug(f"input data is {input_data}, shape is {image.size}")
        logger.info("step 1: Preprocess successfully")
        return input_source, output_methods

    @Timer.time_recorder
    def postprocess(self, data, output_methods):
        result = parse(data, output_methods)
        return result

    @Timer.time_recorder
    def inference(self, data):
        result = self.yolov5_model(data)
        logger.info("step 2: Inference successfully")
        return result

    @Timer.time_recorder
    def handle(self, data: dict, context):
        """Entry point for default handler. It takes the data from the input request and returns
           the predicted outcome for the input.

        Args:
            data (dict): The input data that needs to be made a prediction request on.
            context (Context): It is a JSON Object containing information pertaining to
                               the model artefacts parameters.

        Returns:
            list : Returns a list of dictionary with the predicted response.
        """
        self.context = context
        input_source, output_methods = self.preprocess(data)
        prediction = self.inference(input_source)
        output = self.postprocess(prediction, output_methods)
        return output