| | 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"))) |
| |
|
