PaddleClas/deploy/paddleserving/recognition/test_cpp_serving_client.py

# 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 sys
import numpy as np

from paddle_serving_client import Client
from paddle_serving_app.reader import *
import cv2
import faiss
import os
import pickle


class MainbodyDetect():
    """
    pp-shitu mainbody detect.
    include preprocess, process, postprocess
    return detect results
    Attention: Postprocess include num limit and box filter; no nms 
    """

    def __init__(self):
        self.preprocess = DetectionSequential([
            DetectionFile2Image(), DetectionNormalize(
                [0.485, 0.456, 0.406], [0.229, 0.224, 0.225], True),
            DetectionResize(
                (640, 640), False, interpolation=2), DetectionTranspose(
                    (2, 0, 1))
        ])

        self.client = Client()
        self.client.load_client_config(
            "picodet_PPLCNet_x2_5_mainbody_lite_v1.0_client/serving_client_conf.prototxt"
        )
        self.client.connect(['127.0.0.1:9293'])

        self.max_det_result = 5
        self.conf_threshold = 0.2

    def predict(self, imgpath):
        im, im_info = self.preprocess(sys.argv[1])
        im_shape = np.array(im.shape[1:]).reshape(-1)
        scale_factor = np.array(list(im_info['scale_factor'])).reshape(-1)
        fetch_map = self.client.predict(
            feed={
                "image": im,
                "im_shape": im_shape,
                "scale_factor": scale_factor,
            },
            fetch=["save_infer_model/scale_0.tmp_1"],
            batch=False)
        return self.postprocess(fetch_map, imgpath)

    def postprocess(self, fetch_map, imgpath):
        #1. get top max_det_result
        det_results = fetch_map["save_infer_model/scale_0.tmp_1"]
        if len(det_results) > self.max_det_result:
            boxes_reserved = fetch_map[
                "save_infer_model/scale_0.tmp_1"][:self.max_det_result]
        else:
            boxes_reserved = det_results

        #2. do conf threshold
        boxes_list = []
        for i in range(boxes_reserved.shape[0]):
            if (boxes_reserved[i, 1]) > self.conf_threshold:
                boxes_list.append(boxes_reserved[i, :])

        #3. add origin image box
        origin_img = cv2.imread(imgpath)
        boxes_list.append(
            np.array([0, 1.0, 0, 0, origin_img.shape[1], origin_img.shape[0]]))
        return np.array(boxes_list)


class ObjectRecognition():
    """
    pp-shitu object recognion for all objects detected by MainbodyDetect.
    include preprocess, process, postprocess
    preprocess include preprocess for each image and batching.
    Batch process
    postprocess include retrieval and nms
    """

    def __init__(self):
        self.client = Client()
        self.client.load_client_config(
            "general_PPLCNet_x2_5_lite_v1.0_client/serving_client_conf.prototxt"
        )
        self.client.connect(["127.0.0.1:9294"])

        self.seq = Sequential([
            BGR2RGB(), Resize((224, 224)), Div(255),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225],
                      False), Transpose((2, 0, 1))
        ])

        self.searcher, self.id_map = self.init_index()

        self.rec_nms_thresold = 0.05
        self.rec_score_thres = 0.5
        self.feature_normalize = True
        self.return_k = 1

    def init_index(self):
        index_dir = "../../drink_dataset_v1.0/index"
        assert os.path.exists(os.path.join(
            index_dir, "vector.index")), "vector.index not found ..."
        assert os.path.exists(os.path.join(
            index_dir, "id_map.pkl")), "id_map.pkl not found ... "

        searcher = faiss.read_index(os.path.join(index_dir, "vector.index"))

        with open(os.path.join(index_dir, "id_map.pkl"), "rb") as fd:
            id_map = pickle.load(fd)
        return searcher, id_map

    def predict(self, det_boxes, imgpath):
        #1. preprocess
        batch_imgs = []
        origin_img = cv2.imread(imgpath)
        for i in range(det_boxes.shape[0]):
            box = det_boxes[i]
            x1, y1, x2, y2 = [int(x) for x in box[2:]]
            cropped_img = origin_img[y1:y2, x1:x2, :].copy()
            tmp = self.seq(cropped_img)
            batch_imgs.append(tmp)
        batch_imgs = np.array(batch_imgs)

        #2. process
        fetch_map = self.client.predict(
            feed={"x": batch_imgs}, fetch=["features"], batch=True)
        batch_features = fetch_map["features"]

        #3. postprocess
        if self.feature_normalize:
            feas_norm = np.sqrt(
                np.sum(np.square(batch_features), axis=1, keepdims=True))
            batch_features = np.divide(batch_features, feas_norm)
        scores, docs = self.searcher.search(batch_features, self.return_k)

        results = []
        for i in range(scores.shape[0]):
            pred = {}
            if scores[i][0] >= self.rec_score_thres:
                pred["bbox"] = [int(x) for x in det_boxes[i, 2:]]
                pred["rec_docs"] = self.id_map[docs[i][0]].split()[1]
                pred["rec_scores"] = scores[i][0]
                results.append(pred)
        return self.nms_to_rec_results(results)

    def nms_to_rec_results(self, results):
        filtered_results = []
        x1 = np.array([r["bbox"][0] for r in results]).astype("float32")
        y1 = np.array([r["bbox"][1] for r in results]).astype("float32")
        x2 = np.array([r["bbox"][2] for r in results]).astype("float32")
        y2 = np.array([r["bbox"][3] for r in results]).astype("float32")
        scores = np.array([r["rec_scores"] for r in results])

        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]
        while order.size > 0:
            i = order[0]
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            inds = np.where(ovr <= self.rec_nms_thresold)[0]
            order = order[inds + 1]
            filtered_results.append(results[i])
        return filtered_results


if __name__ == "__main__":
    det = MainbodyDetect()
    rec = ObjectRecognition()

    #1. get det_results
    det_results = det.predict(sys.argv[1])
    print(det_results)

    #2. get rec_results
    rec_results = rec.predict(det_results, sys.argv[1])
    print(rec_results)