youtureid.py

Go to the documentation of this file.

# This file is part of OpenCV Zoo project.
# It is subject to the license terms in the LICENSE file found in the same directory.
#
# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
# Third party copyrights are property of their respective owners.
 
import numpy as np
import cv2 as cv
 
class YoutuReID:
    def __init__(self, modelPath, backendId=0, targetId=0):
        self._modelPath = modelPath
        self._backendId = backendId
        self._targetId = targetId
 
        self._model = cv.dnn.readNet(modelPath)
        self._model.setPreferableBackend(self._backendId)
        self._model.setPreferableTarget(self._targetId)
 
        self._input_size = (128, 256) # fixed
        self._output_dim = 768
        self._mean = (0.485, 0.456, 0.406)
        self._std = (0.229, 0.224, 0.225)
 
    @property
    def name(self):
        return self.__class__.__name__
 
    def setBackend(self, backend_id):
        self._backendId = backend_id
        self._model.setPreferableBackend(self._backendId)
 
    def setTarget(self, target_id):
        self._targetId = target_id
        self._model.setPreferableTarget(self._targetId)
 
    def _preprocess(self, image):
        image = image[:, :, ::-1]
        image = (image / 255.0 - self._mean) / self._std
        return cv.dnn.blobFromImage(image.astype(np.float32))
        # return cv.dnn.blobFromImage(image, scalefactor=(1.0/255.0), size=self._input_size, mean=self._mean) / self._std
 
    def infer(self, image):
        # Preprocess
        inputBlob = self._preprocess(image)
 
        # Forward
        self._model.setInput(inputBlob)
        features = self._model.forward()
        return np.reshape(features, (features.shape[0], features.shape[1]))
 
    def query(self, query_img_list, gallery_img_list, topK=5):
        query_features_list = []
        for q in query_img_list:
            query_features_list.append(self.infer(q))
        query_features = np.concatenate(query_features_list, axis=0)
        query_norm = np.linalg.norm(query_features, ord=2, axis=1, keepdims=True)
        query_arr = query_features / (query_norm + np.finfo(np.float32).eps)
 
        gallery_features_list = []
        for g in gallery_img_list:
            gallery_features_list.append(self.infer(g))
        gallery_features = np.concatenate(gallery_features_list, axis=0)
        gallery_norm = np.linalg.norm(gallery_features, ord=2, axis=1, keepdims=True)
        gallery_arr = gallery_features / (gallery_norm + np.finfo(np.float32).eps)
 
        dist = np.matmul(query_arr, gallery_arr.T)
        idx = np.argsort(-dist, axis=1)
        return [i[0:topK] for i in idx]