import torch
from plyfile import PlyData
from torch_geometric.data import Data


def read_ply(path):
    with open(path, 'rb') as f:
        data = PlyData.read(f)

    pos = ([torch.tensor(data['vertex'][axis]) for axis in ['x', 'y', 'z']])
    pos = torch.stack(pos, dim=-1)

    face = None
    if 'face' in data:
        faces = data['face']['vertex_indices']
        faces = [torch.tensor(face, dtype=torch.long) for face in faces]
        face = torch.stack(faces, dim=-1)

    data = Data(pos=pos)
    data.face = face

    return data

[(11, 'torch.stack', 'torch.stack', False, 'import torch\n'), (19, 'torch_geometric.data.Data', 'Data', False, 'from torch_geometric.data import Data\n'), (8, 'plyfile.PlyData.read', 'PlyData.read', False, 'from plyfile import PlyData\n'), (10, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (17, 'torch.stack', 'torch.stack', False, 'import torch\n'), (16, 'torch.tensor', 'torch.tensor', False, 'import torch\n')]