reid_demo.py

Go to the documentation of this file.

import sys
 
sys_path = ['/workspace/sm_action', '/workspace', '/workspace/smreid']
for path in sys_path:
    if path not in sys.path:
        sys.path.append(path)
 
import os
import cv2
import numpy as np
from tqdm import tqdm
import time
import torch
 
from reid_api import *
 
 
from fastreid.engine import DefaultPredictor
from smutils.utils_os import search_file
 
if __name__ == "__main__":
    
    #모델 설정 값
    device = 'cuda:1'
    cfg_path = '/workspace/configs/reid/mgn/config.yaml'
    checkpoint_path = '/workspace/checkpoints/reid/mgn/market_mgn_R50-ibn.pth'
    crop_image_folder = '/workspace/a_track_result_crop' # 테스트할 crop 이미지들이 있는 폴더
 
    #모델 빌드
    model = build_reid_model(cfg_path, checkpoint_path, device)
    
    #이미지 폴더에 있는 이미지 리스트
    file_list, file_path = search_file(crop_image_folder, '.jpg')
 
    #1000개 이하로 선택
    file_path = file_path[0:min(len(file_path), 1000)]
 
    feats = get_reid_feats(model, file_path) #특징 생성
    distmat = get_reid_matrix(feats) #유사도 매트릭스 생성
 
    #1번 인덱스의 특징과 비교했을때 유사도 0.9 이상인 특징 찾기
    s_sim, keep_index = get_reid_image_index(distmat, select_id= 1, threshold=0.9) 
    sorted_keep_index = index_sort(keep_index, s_sim) #내림 차순 정렬
 
    #출력
    for idx, sim in sorted_keep_index:
        print(f'{file_list[idx]}, {sim}')