reid_api.py

Go to the documentation of this file.

import os
import cv2
import operator
 
import numpy as np
from tqdm import tqdm
 
import torch
from fastreid.engine import DefaultPredictor
from fastreid.config import get_cfg
 
# reid cfg 만들기
def make_reid_cfg(cfg_path, checkpoint_path, device='cuda:0'):
    """
    fast-reid 모델 빌드를 위한 cfg 생성 기능
    args:
        cfg_path (str): 모델이 정의된 config 파일 경로
        checkpoint_path (str): 모델의 파라미터가 저장된 경로
        device (str): 모델이 구동될 디바이스
    return: cfg
    """
    cfg = get_cfg()
    cfg.merge_from_file(cfg_path)
    cfg.merge_from_list(['MODEL.DEVICE', device])
    cfg.merge_from_list(['MODEL.WEIGHTS', checkpoint_path])
    
    return cfg
 
# reid model build
def build_reid_model(args):
    """
    fast-reid 모델 빌드 기능
    args:
        cfg_path (str): 모델이 정의된 config 파일 경로
        checkpoint_path (str): 모델의 파라미터가 저장된 경로
        device (str): 모델이 구동될 디바이스
    return: reid 모델
    """
    reid_cfg = make_reid_cfg(args.cfg_path, args.checkpoint_path, args.device)
    model = DefaultPredictor(reid_cfg)
 
    return model
 
# image : BGR type numpy array, using opecv read
def inference_reid_model(model, image):
    """
    이미지 한장에 대해 reid 모델을 inference 하는 기능
    args:
        model: reid 모델
        image (np.array): BGR 이미지
    return:
        reid features
    """
    image = image[:, :, ::-1] #RGB 변환
    image = cv2.resize(image, tuple(model.cfg.INPUT.SIZE_TEST[::-1]), interpolation=cv2.INTER_CUBIC) #리사이즈
    image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))[None] #텐서 변환
 
    #구동
    return model(image)
 
# 이미지 path들에서 이미지를 불러와 reid 모델에 넣고, 해당 값들을 feats에 저장
def get_reid_feats(model, crop_images):
    """
    여러장의 이미지에 대해 reid 특징을 추출하는 기능
    args:
        model: reid 모델
        crop_images (list[str or np.array]): 사람 영역만 자른 이미지 저장 경로 또는 이미지
    return (list[Tensor]):
        각 이미지에 대한 reid 특징
    """
    feats = []
 
    for img in crop_images:
        #입력이 이미지 경로일 경우 로드
        if isinstance(img, str):
            img = img = cv2.imread(img)
 
        # model inference
        inference_result = inference_reid_model(model, img)
        feats.append(inference_result)
    
    return feats
 
# 쿼리와 갤러기 간의 유사도 행렬 계산
# def get_reid_matrix(feats):
#     if not isinstance(feats, type(torch.float32)):
#         feats = torch.cat(feats, dim=0)
 
#     # compute cosine distance
#     simmat = torch.mm(feats, feats.t())
#     simmat = simmat.numpy()    
#     return simmat
 
def get_reid_matrix(q_feats, g_feats=None):
    """
    쿼리와 갤러리 간의 유사도 행렬을 계산하는 기능
    args:
        q_feats (list[tensor] or Tensor): 쿼리, reid 특징 벡터
        g_feats (Tensor): 갤러리, reid 특징 벡터
    return (np.array):
        유사도 행렬
    """
    # q_feats, g_feats ( list[tensor]   : tensor shape [1, 2048 or any ])
    #                       tensor      : tensor shape [N, 2048 or any]
    if isinstance(q_feats, list):
        q_feats = torch.cat(q_feats, dim=0)
 
    #갤러리의 입력이 없으면 쿼리원소들 간의 유사도 측정
    if g_feats == None:
        g_feats = q_feats
 
    if isinstance(g_feats, list):
        g_feats = torch.cat(g_feats, dim=0)
 
    # compute cosine distance
    simmat = torch.mm(q_feats, g_feats.t())
    simmat = simmat.numpy()    
    return simmat
 
# 찾고자 하는 id와 넘어야 하는 유사도로 원하는 feature 찾기
def get_reid_image_index(simmat, select_id = 1, threshold=0.8):
    """
    유사도 매트릭스를 기반으로 임계치 이상의 유사도를 가지는 feature를 찾는 기능
    args:
        simmat (np.array): 유사도 매트릭스
        select_id (int): 유사도를 비교하려는 인덱스
        threshold (float): 유사도 임계치
    return:
        s_sim (np.array):select_id와 다른 특징과의 유사도
        keep_index (np.array): 임계치 이상의 유사도를 가지는 인덱스
    """
    # select_id에 대한 유사도
    s_sim = simmat[select_id, :]
    
    # 임계값을 넘는 feat들만 받기
    keep = s_sim > threshold
    keep_index = np.where(keep)[0]
    
    return s_sim, keep_index
 
 
# 유사도 score에 따라 정렬
def index_sort(keep_index, s_sim):
    """
    유사도를 기준으로 정렬하는 기능
    args:
        keep_index (np.array): 인덱스
        s_sim (np.array): 유사도
    return (list[tuple]): (인덱스, 유사도) 튜플의 내림차순 정렬된 리스트
        
    """
    sort_index = []
 
    #(인덱스, 유사도)
    for i in range(len(keep_index)):
        index = keep_index[i]
        sim = s_sim[index]
        sort_index.append((index,sim))
 
    #내림차순 정렬
    sorted_keep_index = sorted(sort_index, key=operator.itemgetter(1), reverse=True)
    
    return sorted_keep_index