Knowhow/Python, Pytorch

pytorch3d.io.IO 느린 로딩 속도 개선 방법

침닦는수건 2024. 9. 20. 11:04
반응형

pytorch3d.io.IO를 이용해 obj 파일을 읽어서 사용한다. 

from pytorhc3d.io import IO
mesh = IO().load_mesh(MESH_PATH).cuda() # () essential

 

근데 이상하게도 pytorch3d 내장 함수로 파일을 읽으면 현저히 느리다. 약 10MB 정도 크기의 파일을 기준으로 했을 때 대략 trimesh 읽는 속도보다 2배 느리다. 

 

내부적으로 뭘 더 읽는지는 모르겠으나, 대부분은 vertex, face, color 3개만 있으면 되기 때문에 굳이 필요없는 연산일 것 같다. 

 

개선법

 from pytorch3d.structures import Meshes
 from pytorch3d.renderer import TexturesVertex
 
 mesh = trimesh.load_mesh(mesh_path)
 verts = np.asarray(mesh.vertices).astype(np.float32)
 faces = np.asarray(mesh.faces).astype(np.float32)
 textures = mesh.visual.to_color()
 textures = np.asarray(textures.vertex_colors)[:, :3].astype(np.float32)/ 255.0

 verts = torch.from_numpy(verts).cuda()
 faces = torch.from_numpy(faces).cuda()
 textures = torch.from_numpy(textures).cuda()
 mesh = Meshes(verts=[verts], faces=[faces])
 mesh.textures = TexturesVertex(textures.reshape(1, -1, 3))

 

위 방식대로 trimesh로 먼저 읽고 내부 값을 pytorch3d Mesh의 초기값으로 넣어주는 식으로 바꿔쓰면 훨씬 빠르다. 

 

대략 2.93초 걸리던 로딩 시간이 1.4초로 줄어든다. 

 

위 예시에서는 color를 face color로 안넣고 vertex color로 넣었는데 texture가 그리 중요하지 않은 경우에는 이정도로 충분하다. 

반응형