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