| import os |
| import shutil |
| import numpy as np |
| import struct |
| import pickle |
| from scipy.sparse import coo_matrix |
|
|
| def load_flow(filename): |
| |
| assert os.path.isfile(filename), "File not found: {}".format(filename) |
|
|
| flow = None |
| with open(filename, 'rb') as fin: |
| width = struct.unpack('I', fin.read(4))[0] |
| height = struct.unpack('I', fin.read(4))[0] |
| channels = struct.unpack('I', fin.read(4))[0] |
| n_elems = height * width * channels |
|
|
| flow = struct.unpack('f' * n_elems, fin.read(n_elems * 4)) |
| flow = np.asarray(flow, dtype=np.float32).reshape([channels, height, width]) |
|
|
| return flow |
|
|
| def load_graph_info(filename, max_edges, max_nodes): |
|
|
| assert os.path.isfile(filename), "File not found: {}".format(filename) |
|
|
| with open(filename, 'rb') as fin: |
| edge_total_size = struct.unpack('I', fin.read(4))[0] |
| edges = struct.unpack('I' * (int(edge_total_size / 4)), fin.read(edge_total_size)) |
| edges = np.asarray(edges, dtype=np.int16).reshape(-1, 2).transpose() |
| nodes_total_size = struct.unpack('I', fin.read(4))[0] |
| nodes_ids = struct.unpack('I' * (int(nodes_total_size / 4)), fin.read(nodes_total_size)) |
| nodes_ids = np.asarray(nodes_ids, dtype=np.int32).reshape(-1) |
| nodes_ids = np.sort(nodes_ids) |
|
|
| edges_extent = np.zeros((2, max_edges), dtype=np.int16) |
| edges_mask = np.zeros((max_edges), dtype=np.bool) |
| edges_mask[:edges.shape[1]] = 1 |
| edges_extent[:, :edges.shape[1]] = edges |
|
|
| nodes_extent = np.zeros((max_nodes), dtype=np.int32) |
| nodes_mask = np.zeros((max_nodes), dtype=np.bool) |
| nodes_mask[:nodes_ids.shape[0]] = 1 |
| nodes_extent[:nodes_ids.shape[0]] = nodes_ids |
|
|
| fx = struct.unpack('f', fin.read(4))[0] |
| fy = struct.unpack('f', fin.read(4))[0] |
| ox = struct.unpack('f', fin.read(4))[0] |
| oy = struct.unpack('f', fin.read(4))[0] |
| |
| return edges_extent, edges_mask, nodes_extent, nodes_mask, fx, fy, ox, oy |
|
|
| def load_adja_id_info(filename, src_mask, H, W, num_adja, num_neigb): |
|
|
| assert os.path.isfile(filename), "File not found: {}".format(filename) |
| assert num_adja<=8, "Num of adja is larger than 8" |
| assert num_neigb<=8, "Num of neighb is larger than 8" |
| src_v_id = np.zeros((H*W, num_adja), dtype=np.int16) |
| src_neigb_id = np.zeros((H*W, num_neigb), dtype=np.int32) |
| with open(filename, 'rb') as fin: |
| neigb_id, value_id = pickle.load(fin) |
| assert((src_mask.sum())==value_id.shape[0]) |
|
|
| for i in range(num_adja): |
| src_v_id[src_mask.reshape(-1), i] = value_id[:, i] |
| for i in range(num_neigb): |
| src_neigb_id[src_mask.reshape(-1), i] = neigb_id[:, i] |
| src_v_id = src_v_id.transpose().reshape(num_adja, H, W) |
| src_neigb_id = src_neigb_id.transpose().reshape(num_neigb, H, W) |
|
|
| return src_v_id, src_neigb_id |
|
|
| def save_flow(filename, flow_input): |
| flow = np.copy(flow_input) |
|
|
| |
| assert len(flow.shape) == 3 |
| |
| with open(filename, 'wb') as fout: |
| fout.write(struct.pack('I', flow.shape[2])) |
| fout.write(struct.pack('I', flow.shape[1])) |
| fout.write(struct.pack('I', flow.shape[0])) |
| fout.write(struct.pack('={}f'.format(flow.size), *flow.flatten("C"))) |
|
|
