0017-fastAPI调用yoloV3模型

「爱情、让人受尽委屈。」 2023-01-08 02:23 163阅读 0赞

# object\_detection #

import cv2
    import sys
    import numpy as np
    
    class object_detector:
    
        def __init__(self, model, cfg, classes, inputSize=416):
            self.model = model
            self.cfg = cfg
            self.classes = classes
            self.inputSize = inputSize
            self.framework = None
            self.load_model()
    
        def load_model(self):
            if self.model.endswith('weights') and self.cfg.endswith('cfg') and self.classes.endswith('names'):
                self.net = cv2.dnn.readNetFromDarknet(self.cfg, self.model)
                self.LABELS = open(self.classes).read().strip().split("\n")
                self.framework = 'Darknet'
            else:
                sys.exit('Wrong input for model weights and cfg')
            #这里使用CUDA加速，如果没有GPU支持会自动转到CPU推理
            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
    
        def predict(self, srcImg):
            boxes = []
            confidences = []
            classIDs = []
            json_result = {}
            dict_list = []
            (H, W) = srcImg.shape[:2]
    
            # Create a 4D blob from a srcImg.
            if self.framework == 'Darknet':
                blob = cv2.dnn.blobFromImage(srcImg, 1/255.0, (self.inputSize, self.inputSize), swapRB=True, crop=False)
            else:
                print("Not a Darknet models")
    
            ln = self.net.getLayerNames()
            ln = [ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
            # Run a model
            self.net.setInput(blob)
            layerOutputs = self.net.forward(ln)
    
            # 在每层输出上循环
            for output in layerOutputs:
                # 对每个检测进行循环
                for detection in output:
                    scores = detection[5:]
                    classID = np.argmax(scores)
                    confidence = scores[classID]
                    # 过滤掉那些置信度较小的检测结果
                    if confidence > 0.02:
                        # 框后接框的宽度和高度
                        box = detection[0:4] * np.array([W, H, W, H])
                        (centerX, centerY, width, height) = box.astype("int")
                        # 边框的左上角
                        x = int(centerX - (width / 2))
                        y = int(centerY - (height / 2))
                        # 更新检测出来的框
                        boxes.append([x, y, int(width), int(height)])
                        confidences.append(float(confidence))
                        classIDs.append(classID)
    
            # 极大值抑制
            idxs = cv2.dnn.NMSBoxes(boxes, confidences, score_threshold=.4, nms_threshold=.3)
    
            if len(idxs) > 0:
                for i in idxs.flatten():
                    (x, y) = (boxes[i][0], boxes[i][1])
                    (w, h) = (boxes[i][2], boxes[i][3])
    
                    data_location = {}
                    data_result = {}
                    data_det = {}
                    data_location["score"] = float(confidences[i])
                    data_location["left"] = int(x)
                    data_location["top"] = int(y)
                    data_location["width"] = int(w)
                    data_location["height"] = int(h)
                    data_result["score"] = float(confidences[i])
                    data_result["label"] = str(self.LABELS[classIDs[i]])
                    data_det["location"] = data_location
                    data_det["result"] = data_result
                    dict_list.append(data_det)
    
            json_result["code"] = 200
            json_result["data"] = dict_list
            # print(json_result)
            return json_result

# main函数 #

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/18 17:04
    # @Author  : Johnson
    from fastapi import FastAPI
    from pydantic import BaseModel
    import uvicorn
    import cv2
    import random
    import os
    import base64
    import json
    import time
    import numpy as np
    from enum import Enum
    
    # 加载前面的检测类
    from object_detection import object_detector
    
    app = FastAPI()
    
    
    def base64toCv(base64_src):
        img_b64decode = base64.b64decode(base64_src)  # base64解码
        img_array = np.fromstring(img_b64decode, np.uint8)  # 转换np序列
        img_cv = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)  # 转换OpenCV格式
        return img_cv
    
    
    # 将识别的类别加入枚举
    class Targets(str, Enum):
        det_target1 = "face"
    
    
    # 定义接收数据的结构
    class Item(BaseModel):
        base64: str = None  # 图片base64
        target: Targets = None  # 识别类型
    
    
    @app.post('/detector')
    async def calculate(request_data: Item):
        Target = request_data.target
        img_base64 = request_data.base64
    
        # 输出检测的信息和调用时间
        print("Detection for", Target.value, "! Time:", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    
        if Target == "face":
            src_img = base64toCv(img_base64)
            result = face_detector.predict(src_img)
            return result
        else:
            print("Target parameter error!")
            return {"code": 201, "data": []}
    
    
    if __name__ == '__main__':
        # 实例化一个检测器，输入参数分别是weights、cfg、names文件的路径，和模型的输入尺寸。
        face_detector = object_detector("./models/face_model/face.weights",
                                        "./models/face_model/face.cfg",
                                        "./models/face_model/face.names",
                                        416)
        print('Loaded face model!')
        print("Service start!")
        uvicorn.run(app=app,
                    host="0.0.0.0",  # 服务器填写0.0.0.0
                    port=12455,
                    workers=1)

# client #

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/18 17:14
    # @Author  : Johnson
    import requests
    import json
    import base64
    
    
    def test(img_path):
        with open(img_path, 'rb') as f:
            base64_data = base64.b64encode(f.read())
            img = base64_data.decode()
        datas = json.dumps({'base64': img, 'target': "face"})
        rec = requests.post("http://0.0.0.0:12455/detector", data=datas)
        return rec.text
    
    
    result = test('1.jpg')
    print(result)