generate_image_features.py

Go to the documentation of this file.

import numpy as np
import mmcv
import copy
import cv2
import os
from tqdm import tqdm
 
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import timm
 
from smdataset.AnnotationDataManager import AnnotationDataManager
from smutils.utils_os import search_file, create_directory
from smutils.utils_image import crop_image
from smutils.utils_data import save_pkl_data
import smrunner
 
def inference_one_frame(model, image, bboxes, device):
    """
    한프레임에서 검출된 박스에대해 이미지 특징 생성 모델을 inference하는 기능
    args:
        model: pytorch 모델
        image (np.array): 원본 이미지
        bboxes (np.array): 검출 박스(x,y,x,y), shape (num_person, 4 or 5 or 6)
        device (str): 모델이 구동될 디바이스
    return (Tensor): 이미지 특징, shape  (num_person, C)
    """
 
    #이미지 전처리
    preprocess = transforms.Compose([
        transforms.ToPILImage(), #구조변환, np.array -> PILImage
        transforms.Resize((224, 224)), #리사이즈
        transforms.ToTensor(), #텐서로 변환
        transforms.Normalize(mean=[0.4815, 0.4578, 0.4082], std=[0.2686, 0.2613, 0.2758]), #가우시안 정규화
    ])
 
    crop_images = crop_image(image, bboxes) #박스영역 잘라내기
    
    #이미지 전처리
    transformed_images = []
    for img in crop_images:
        transformed_images.append(preprocess(img))
    
    image_batch = torch.stack(transformed_images) #모델 입력 데이터 생성, 모델 inference를 위한 배치단위 변환
    image_batch = image_batch.to(device) #입력 데이터를 디바이스로 전송
 
    #inference
    with torch.no_grad():
        results = model(image_batch) #모델 구동
        results = results.cpu() #결과(이미지 특징) cpu로 전송
    
    return results
 
 
def gen_image_features(model, data_manager, save_path, device):
    """
    하나의 어노테이션 파일을 기반으로 사람영역에 대해 이미지 특징 생성 모델을 사용하여 예측한 이미지 특징을 저장하는 기능
    args:
        model: pytorch 모델
        data_manager: AnnotationDataManager 객체, 어노테이션 데이터파일로 초기화까지 진행
        save_path (str): 이미지 특징을 저장할 경로
        device (str): 모델이 구동될 디바이스
    """
 
    #이미지 숫자 체크
    img_num = data_manager.get_number_of_image()
    image_feats_data = {}
    for i in tqdm(range(img_num)):
        anno_data = data_manager.get_anno_data_in_image() #이미지에 있는 어노테이션 데이터를 가져옴
        img_path = data_manager.get_image_path() #이미지 경로를 가져옴
        
        #어노테이션 박스 구조 변경 (x, y, w, h) -> (x, y, x, y)
        bbox_in_image = []
        for anno in anno_data:
            bbox_tmp = copy.deepcopy(anno['bbox'])
 
            #마진 추가
            margin_x = bbox_tmp[2]*0.15
            margin_y = bbox_tmp[3]*0.15
            
            bbox_tmp[0] -= margin_x
            bbox_tmp[1] -= margin_y
            bbox_tmp[2] *= 1.3
            bbox_tmp[3] *= 1.3
 
            #박스 구조 변경            
            bbox_tmp[2] += bbox_tmp[0] #TODO: 영상 크기로 최대치 체크
            bbox_tmp[3] += bbox_tmp[1] #TODO: 영상 크기로 최대치 체크
            bbox_tmp[0] = max(bbox_tmp[0], 0)
            bbox_tmp[1] = max(bbox_tmp[1], 0)
 
            bbox_in_image.append(bbox_tmp)
 
        #박스가 없으면 다음 이미지로 넘어감
        if len(bbox_in_image) == 0:
            data_manager.move_image_right()
            continue
 
        bboxes = np.array(bbox_in_image) #박스 타입 변경
        image = cv2.imread(img_path) #이미지 로드
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
        #이미지 특징 생성
        image_feats = inference_one_frame(model, image, bboxes, device)
        
        #이미지 특징을 추적 아이디 별로 저장
        for anno, feat in zip(anno_data, image_feats):
            anno_id = anno['id']
            image_id = anno['image_id']
            track_id = anno['track_id']
 
            if track_id not in image_feats_data:
                image_feats_data[track_id] = {}
            
            image_feats_data[track_id][image_id] = feat
 
        #다음 이미지로 이동
        data_manager.move_image_right()
 
    #이미지 특징 저장
    save_pkl_data(image_feats_data, save_path)
 
def run_gen_image_features(args):
    """
    어노테이션 파일들을 기반으로 사람영역에 대해 이미지 특징 생성 모델을 사용하여 예측한 이미지 특징을 저장하는 기능
    이미지 특징 생성 모델은 timm 패키지에서 제공하고 있는 eva02_base_patch16_clip_224 모델(사전학습 파라미터 이용)을 사용함.
    args:
        save_root (str): 이미지 특징을 저장할 폴더
        device (str): 모델이 구동될 디바이스
        dataset_folder_list (list[str]): 어노테이션 파일이 저장된 폴더 리스트
    """
    print('gen image features start')
    #json 저장 폴더 생성
    
    #이미지 특징이 저장될 폴더 생성
    create_directory(args.save_root)
 
    #모델이 구동될 디바이스 설정
    device = args.device
 
    #모델 생성
    model = timm.create_model('eva02_base_patch16_clip_224', pretrained=True) #timm 모델
    model.head = nn.Identity() #헤더를 제거
    model.to(device) #모델을 디바이스로 전송
    model.eval() #eval 모드 전환
 
    #어노테이션 관리자 생성
    data_manager = AnnotationDataManager()
    
    #어노테이션 리스트
    json_name_list = []
    json_path_list = []
    for dataset_folder in args.dataset_folder_list:
        name_list, path_list = search_file(dataset_folder, '.json')
        json_name_list.extend(name_list)
        json_path_list.extend(path_list)
    
    file_num = len(json_path_list) #파일 수량, 진행 정도 표시를 위함
    for i, (json_name, json_path) in enumerate(zip(json_name_list, json_path_list)):
        
        #저장 경로 설정
        save_path = os.path.join(args.save_root, json_name)
        save_path = save_path.replace('.json', '.pkl')
        
        #어노테이션 관리자 초기화
        data_manager.load_annotation(json_path)
        data_manager.init_annotation()
        
        #진행도 출력
        print(f"[ {i+1} / {file_num}] : load path : {json_path}")
        print(f"[ {i+1} / {file_num}] : save path : {save_path}")
 
        #어노테이션 파일 하나에 대해서 이미지 특징 생성 진행
        gen_image_features(model, data_manager, save_path, device)
        print('')