Shortcuts

Source code for mmflow.datasets.utils.flow_io

# Copyright (c) OpenMMLab. All rights reserved.

import re
from io import BytesIO
from typing import Tuple

import cv2
import matplotlib.pyplot as plt
import mmcv
import numpy as np
from numpy import ndarray


[docs]def read_flow(name: str) -> np.ndarray: """Read flow file with the suffix '.flo'. This function is modified from https://lmb.informatik.uni-freiburg.de/resources/datasets/IO.py Copyright (c) 2011, LMB, University of Freiburg. Args: name (str): Optical flow file path. Returns: ndarray: Optical flow """ with open(name, 'rb') as f: header = f.read(4) if header.decode('utf-8') != 'PIEH': raise Exception('Flow file header does not contain PIEH') width = np.fromfile(f, np.int32, 1).squeeze() height = np.fromfile(f, np.int32, 1).squeeze() flow = np.fromfile(f, np.float32, width * height * 2).reshape( (height, width, 2)) return flow
[docs]def write_flow(flow: np.ndarray, flow_file: str) -> None: """Write the flow in disk. This function is modified from https://lmb.informatik.uni-freiburg.de/resources/datasets/IO.py Copyright (c) 2011, LMB, University of Freiburg. Args: flow (ndarray): The optical flow that will be saved. flow_file (str): The file for saving optical flow. """ with open(flow_file, 'wb') as f: f.write('PIEH'.encode('utf-8')) np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f) flow = flow.astype(np.float32) flow.tofile(f)
[docs]def visualize_flow(flow: np.ndarray, save_file: str = None) -> np.ndarray: """Flow visualization function. Args: flow (ndarray): The flow will be render save_dir ([type], optional): save dir. Defaults to None. Returns: ndarray: flow map image with RGB order. """ # return value from mmcv.flow2rgb is [0, 1.] with type np.float32 flow_map = np.uint8(mmcv.flow2rgb(flow) * 255.) if save_file: plt.imsave(save_file, flow_map) return flow_map
[docs]def render_color_wheel(save_file: str = 'color_wheel.png') -> np.ndarray: """Render color wheel. Args: save_file (str): The saved file name . Defaults to 'color_wheel.png'. Returns: ndarray: color wheel image. """ x0 = 75 y0 = 75 height = 151 width = 151 flow = np.zeros((height, width, 2), dtype=np.float32) grid_x = np.tile(np.expand_dims(np.arange(width), 0), [height, 1]) grid_y = np.tile(np.expand_dims(np.arange(height), 1), [1, width]) grid_x0 = np.tile(np.array([x0]), [height, width]) grid_y0 = np.tile(np.array([y0]), [height, width]) flow[:, :, 0] = grid_x - grid_x0 flow[:, :, 1] = grid_y - grid_y0 return visualize_flow(flow, save_file)
[docs]def read_flow_kitti(name: str) -> Tuple[np.ndarray, np.ndarray]: """Read sparse flow file from KITTI dataset. This function is modified from https://github.com/princeton-vl/RAFT/blob/master/core/utils/frame_utils.py. Copyright (c) 2020, princeton-vl Licensed under the BSD 3-Clause License Args: name (str): The flow file Returns: Tuple[ndarray, ndarray]: flow and valid map """ # to specify not to change the image depth (16bit) flow = cv2.imread(name, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR) flow = flow[:, :, ::-1].astype(np.float32) # flow shape (H, W, 2) valid shape (H, W) flow, valid = flow[:, :, :2], flow[:, :, 2] flow = (flow - 2**15) / 64.0 return flow, valid
[docs]def write_flow_kitti(uv: np.ndarray, filename: str): """Write the flow in disk. This function is modified from https://github.com/princeton-vl/RAFT/blob/master/core/utils/frame_utils.py. Copyright (c) 2020, princeton-vl Licensed under the BSD 3-Clause License Args: uv (ndarray): The optical flow that will be saved. filename ([type]): The file for saving optical flow. """ uv = 64.0 * uv + 2**15 valid = np.ones([uv.shape[0], uv.shape[1], 1]) uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16) cv2.imwrite(filename, uv[..., ::-1])
def flow_from_bytes(content: bytes, suffix: str = 'flo') -> ndarray: """Read dense optical flow from bytes. .. note:: This load optical flow function works for FlyingChairs, FlyingThings3D, Sintel, FlyingChairsOcc datasets, but cannot load the data from ChairsSDHom. Args: content (bytes): Optical flow bytes got from files or other streams. Returns: ndarray: Loaded optical flow with the shape (H, W, 2). """ assert suffix in ('flo', 'pfm'), 'suffix of flow file must be `flo` '\ f'or `pfm`, but got {suffix}' if suffix == 'flo': return flo_from_bytes(content) else: return pfm_from_bytes(content) def flo_from_bytes(content: bytes): """Decode bytes based on flo file. Args: content (bytes): Optical flow bytes got from files or other streams. Returns: ndarray: Loaded optical flow with the shape (H, W, 2). """ # header in first 4 bytes header = content[:4] if header != b'PIEH': raise Exception('Flow file header does not contain PIEH') # width in second 4 bytes width = np.frombuffer(content[4:], np.int32, 1).squeeze() # height in third 4 bytes height = np.frombuffer(content[8:], np.int32, 1).squeeze() # after first 12 bytes, all bytes are flow flow = np.frombuffer(content[12:], np.float32, width * height * 2).reshape( (height, width, 2)) return flow def pfm_from_bytes(content: bytes) -> np.ndarray: """Load the file with the suffix '.pfm'. Args: content (bytes): Optical flow bytes got from files or other streams. Returns: ndarray: The loaded data """ file = BytesIO(content) color = None width = None height = None scale = None endian = None header = file.readline().rstrip() if header == b'PF': color = True elif header == b'Pf': color = False else: raise Exception('Not a PFM file.') dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline()) if dim_match: width, height = list(map(int, dim_match.groups())) else: raise Exception('Malformed PFM header.') scale = float(file.readline().rstrip()) if scale < 0: # little-endian endian = '<' scale = -scale else: endian = '>' # big-endian data = np.frombuffer(file.read(), endian + 'f') shape = (height, width, 3) if color else (height, width) data = np.reshape(data, shape) data = np.flipud(data) return data[:, :, :-1] def read_pfm(file: str) -> np.ndarray: """Load the file with the suffix '.pfm'. This function is modified from https://lmb.informatik.uni-freiburg.de/resources/datasets/IO.py Copyright (c) 2011, LMB, University of Freiburg. Args: file (str): The file name will be loaded Returns: ndarray: The loaded data """ file = open(file, 'rb') color = None width = None height = None scale = None endian = None header = file.readline().rstrip() if header.decode('ascii') == 'PF': color = True elif header.decode('ascii') == 'Pf': color = False else: raise Exception('Not a PFM file.') dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode('ascii')) if dim_match: width, height = list(map(int, dim_match.groups())) else: raise Exception('Malformed PFM header.') scale = float(file.readline().decode('ascii').rstrip()) if scale < 0: # little-endian endian = '<' scale = -scale else: endian = '>' # big-endian data = np.fromfile(file, endian + 'f') shape = (height, width, 3) if color else (height, width) data = np.reshape(data, shape) data = np.flipud(data) return data[:, :, :-1]
Read the Docs v: latest
Versions
latest
stable
Downloads
pdf
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.