generate_image_features Namespace Reference

Functions
	inference_one_frame (model, image, bboxes, device)
	gen_image_features (model, data_manager, save_path, device)
	run_gen_image_features (args)

Function Documentation

gen_image_features()

generate_image_features.gen_image_features	(		model,
			data_manager,
			save_path,
			device )

하나의 어노테이션 파일을 기반으로 사람영역에 대해 이미지 특징 생성 모델을 사용하여 예측한 이미지 특징을 저장하는 기능
args:
    model: pytorch 모델
    data_manager: AnnotationDataManager 객체, 어노테이션 데이터파일로 초기화까지 진행
    save_path (str): 이미지 특징을 저장할 경로
    device (str): 모델이 구동될 디바이스

Definition at line 56 of file generate_image_features.py.

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)
 

inference_one_frame()

generate_image_features.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)

Definition at line 19 of file generate_image_features.py.

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
 
 

run_gen_image_features()

generate_image_features.run_gen_image_features

(

args

)

어노테이션 파일들을 기반으로 사람영역에 대해 이미지 특징 생성 모델을 사용하여 예측한 이미지 특징을 저장하는 기능
이미지 특징 생성 모델은 timm 패키지에서 제공하고 있는 eva02_base_patch16_clip_224 모델(사전학습 파라미터 이용)을 사용함.
args:
    save_root (str): 이미지 특징을 저장할 폴더
    device (str): 모델이 구동될 디바이스
    dataset_folder_list (list[str]): 어노테이션 파일이 저장된 폴더 리스트

Definition at line 124 of file generate_image_features.py.

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('')