Serving 服务编写（Predictor 方式）

import json

class Predictor:
    def __init__(self):
        with open('classes.json', 'r') as f:
            values = json.load(f)
        self.values = values

    ...

import openbayes_serving as serv


class Predictor:
    def __init__(self):
        """
        负责加载相应的模型以及对元数据的初始化
        """
        pass

    def predict(self, json):
        """
        在每次请求都会被调用
        接受 HTTP 请求的内容（`json`)
        进行必要的预处理（preprocess）后预测结果，
        最终将结果进行后处理（postprocess）并返回给调用方

        Args:
            json: 请求的数据

        Returns:
            预测结果
        """
        pass

if __name__ == '__main__':  # 如果直接执行了 predictor.py，而不是被其他文件 import
    serv.run(Predictor)  # 开始提供服务

from io import BytesIO

import requests
import torch
from PIL import Image
from torchvision import models
from torchvision import transforms

import openbayes_serving as serv


class Predictor:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"using device: {self.device}")

        model = models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
        model.load_state_dict(torch.load("fasterrcnn_resnet50_fpn_coco-258fb6c6.pth"))
        model = model.to(self.device)
        model.eval()

        self.preprocess = transforms.Compose([transforms.ToTensor()])

        with open("coco_labels.txt") as f:
            self.coco_labels = f.read().splitlines()

        self.model = model

    def predict(self, json):
        # 只使用了 json 参数，其内容为 Dict 格式，直接通过 json[key] 的形式获取内容即可
        threshold = float(json["threshold"])
        image = requests.get(json["url"]).content
        img_pil = Image.open(BytesIO(image))
        img_tensor = self.preprocess(img_pil).to(self.device)
        img_tensor.unsqueeze_(0)

        with torch.no_grad():
            pred = self.model(img_tensor)

        predicted_class = [self.coco_labels[i] for i in pred[0]["labels"].cpu().tolist()]
        predicted_boxes = [
            [(i[0], i[1]), (i[2], i[3])] for i in pred[0]["boxes"].detach().cpu().tolist()
        ]
        predicted_score = pred[0]["scores"].detach().cpu().tolist()
        predicted_t = [predicted_score.index(x) for x in predicted_score if x > threshold]
        if len(predicted_t) == 0:
            return [], []

        predicted_t = predicted_t[-1]
        predicted_boxes = predicted_boxes[: predicted_t + 1]
        predicted_class = predicted_class[: predicted_t + 1]
        return predicted_boxes, predicted_class


if __name__ == '__main__':
    serv.run(Predictor)

class PredictorExample1:

    def __init__(self):
        pass

    def predict(self, json):
        return {'result': 'It works!'}

class PredictorExample2:

    def __init__(self, onnx):
        self.onnx = onnx

    def preprocess(self, v):
        ...

    def postprocess(self, v):
        ...

    def predict(self, json):
        onnx = self.onnx
        m = self.preprocess(json['data'])
        result = self.postprocess(onnx.run(None, {'data': m}))
        return result

class PredictorExample1:

    def __init__(self):
        pass

    def predict(self, json, params):
        return {'message': f'Param foo is {params.get("foo")}'}

def predict(self, json):
    return [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

def predict(self):
    return "class 1"

def predict(self):
    # bytes 类型的返回
    array = np.random.randn(3, 3)
    response = pickle.dumps(array)
    return response

def predict(self):
    from flask import Response
    data = b"class 1"
    response = Response(data, mimetype="text/plain")
    return response