카테고리 없음

[torch] Arcface torch tensor로 warp_affine 하기

3pie 2024. 3. 14. 17:37

Arcface tnesoRT로 경량화 진행하면서 preprocess를 torch gpu로 연산을해야 속도가 빨라지기때문에 torch로 변환해서 진행해야했다.

 

기존코드의 경우에는 numpy로 진행되는건데 다른걸로 변경 

image = kornia.geometry.transform.warp_affine(img_list, M, (image_size, image_size), mode='bilinear', padding_mode='zeros', align_corners=True, fill_value=torch.zeros(3))

 

from torchgeometry.core.imgwarp import warp_affine

제공하는것 사용

 

warped = warp_affine(input_frame,M,(image_size, image_size),flags='bilinear',padding_mode='zeros')

 

요렇게 지금은 1배치로 처리되게 되어 있는다 배치수를 늘려서 더 빠르게 진행하시길 

input_frame = img.float()
input_frame = [input_frame for _ in range(1)]
input_frame = torch.stack(input_frame)
M, pose_index = self.estimate_norm(landmark, image_size, mode)
M = torch.tensor(M).float()
M = [M for _ in range(1)]
M = torch.stack(M).to(torch.device("cuda"))

warped = warp_affine(input_frame,M,(image_size, image_size),flags='bilinear',padding_mode='zeros')