Deep-JGAC/inout.py at main · SYSU-Video/Deep-JGAC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import open3d as o3d
import os, time
import numpy as np
import h5py

def read_h5(filedir, dtype="int32"):
    pc = h5py.File(filedir, 'r')['data'][:]
    coords = pc[:,0:3].astype(dtype)

    return coords

def write_h5(filedir, coords, dtype="int32"):
    data = coords.astype(dtype)
    with h5py.File(filedir, 'w') as h:
        h.create_dataset('data', data=data, shape=data.shape)

    return

def read_ply_ascii(filedir, dtype="int32"):
    files = open(filedir, 'r')
    data = []
    for i, line in enumerate(files):
        wordslist = line.split(' ')
        try:
            line_values = []
            for i, v in enumerate(wordslist):
                if v == '\n': continue
                line_values.append(float(v))
        except ValueError: continue
        data.append(line_values)
    data = np.array(data)
    coords = data[:,0:3].astype(dtype)

    return coords

def write_ply_ascii(filedir, coords, dtype='int32'):
    if os.path.exists(filedir): os.system('rm '+filedir)
    f = open(filedir,'a+')
    f.writelines(['ply\n','format ascii 1.0\n'])
    f.write('element vertex '+str(coords.shape[0])+'\n')
    f.writelines(['property float x\n','property float y\n','property float z\n'])
    f.write('end_header\n')
    coords = coords.astype(dtype)
    for p in coords:
        f.writelines([str(p[0]), ' ', str(p[1]), ' ',str(p[2]), '\n'])
    f.close()

    return

def read_ply_o3d(filedir, dtype='int32'):
    pcd = o3d.io.read_point_cloud(filedir)
    coords = np.asarray(pcd.points).astype(dtype)
    colors = np.asarray(pcd.colors)
    return coords,colors

def write_ply_o3d(filedir, coords, colors,dtype='int32', normal=False, knn=None,withcolor = True):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coords.astype(dtype))
    pcd.colors = o3d.utility.Vector3dVector(colors)
    if normal:
        assert knn is not None
        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=knn))
    o3d.io.write_point_cloud(filedir, pcd, write_ascii=True)
    f = open(filedir)
    lines = f.readlines()
    lines[4] = 'property float x\n'
    lines[5] = 'property float y\n'
    lines[6] = 'property float z\n'
    if normal:
        lines[7] = 'property float nx\n'
        lines[8] = 'property float ny\n'
        lines[9] = 'property float nz\n'
    # if withcolor:
    #     lines[10] = 'property uchar red\n'
    #     lines[11] = 'property uchar green\n'
    #     lines[12] = 'property uchar blue\n'
    fo = open(filedir, "w")
    fo.writelines(lines)

    return

def read_bin(filedir, dtype="float32"):
    """kitti
    """
    data = np.fromfile(filedir, dtype=dtype).reshape(-1, 4)
    coords = data[:,:3]

    return coords

def read_coords(filedir):
    if filedir.endswith('h5'): coords = read_h5(filedir)
    if filedir.endswith('ply'): coords = read_ply_o3d(filedir)
    if filedir.endswith('bin'): coords = read_bin(filedir)

    return coords