[Features] Load optical flow data from bytes (#1362)

* [Features] Load optical flow data from bytes

* docstring

* revise base on comments

* compression test data

* compression test data

* docstring

* minors

mmcv/video/__init__.py

 # Copyright (c) OpenMMLab. All rights reserved.
 from .io import Cache, VideoReader, frames2video
-from .optflow import (dequantize_flow, flow_warp, flowread, flowwrite,
-                      quantize_flow)
+from .optflow import (dequantize_flow, flow_from_bytes, flow_warp, flowread,
+                      flowwrite, quantize_flow, sparse_flow_from_bytes)
 from .processing import concat_video, convert_video, cut_video, resize_video
 
 __all__ = [
     'Cache', 'VideoReader', 'frames2video', 'convert_video', 'resize_video',
     'cut_video', 'concat_video', 'flowread', 'flowwrite', 'quantize_flow',
-    'dequantize_flow', 'flow_warp'
+    'dequantize_flow', 'flow_warp', 'flow_from_bytes', 'sparse_flow_from_bytes'
 ]

mmcv/video/optflow.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
 
+import cv2
 import numpy as np
 
 from mmcv.arraymisc import dequantize, quantize
@@ -198,3 +199,56 @@ def flow_warp(img, flow, filling_value=0, interpolate_mode='nearest'):
             'We only support interpolation modes of nearest and bilinear, '
             f'but got {interpolate_mode}.')
     return output.astype(img.dtype)
+
+
+def flow_from_bytes(content):
+    """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).
+    """
+
+    # header in first 4 bytes
+    header = content[:4]
+    if header.decode('utf-8') != '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 sparse_flow_from_bytes(content):
+    """Read the optical flow in KITTI datasets from bytes.
+
+    This function is modified from RAFT load the `KITTI datasets
+    <https://github.com/princeton-vl/RAFT/blob/224320502d66c356d88e6c712f38129e60661e80/core/utils/frame_utils.py#L102>`_.
+
+    Args:
+        content (bytes): Optical flow bytes got from files or other streams.
+
+    Returns:
+        Tuple(ndarray, ndarray): Loaded optical flow with the shape (H, W, 2)
+            and flow valid mask with the shape (H, W).
+    """  # nopa
+
+    content = np.frombuffer(content, np.uint8)
+    flow = cv2.imdecode(content, 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

tests/test_video/test_optflow.py

@@ -3,6 +3,7 @@ import os
 import os.path as osp
 import tempfile
 
+import cv2
 import numpy as np
 import pytest
 from numpy.testing import assert_array_almost_equal, assert_array_equal
@@ -245,3 +246,46 @@ def test_make_color_wheel():
                  [1. , 0. , 1. ],  # noqa
                  [1. , 0. , 0.5]], dtype=np.float32))  # noqa
     # yapf: enable
+
+
+def test_flow_from_bytes():
+    data_dir = osp.join(osp.dirname(__file__), '../data')
+    flow_shape = (60, 80, 2)
+    flow_file = osp.join(data_dir, 'optflow.flo')
+
+    # read .flo file
+    flow_fromfile = mmcv.flowread(flow_file)
+
+    with open(flow_file, 'rb') as f:
+        flow_bytes = f.read()
+    flow_frombytes = mmcv.flow_from_bytes(flow_bytes)
+
+    assert flow_frombytes.shape == flow_shape
+    assert np.all(flow_frombytes == flow_fromfile)
+
+
+def test_sparse_flow_from_bytes():
+    data_dir = osp.join(osp.dirname(__file__), '../data')
+    flow_file = osp.join(data_dir, 'sparse_flow.png')
+
+    with open(flow_file, 'rb') as f:
+        flow_bytes = f.read()
+    # read flow from bytes
+    flow_frombytes, valid_frombytes = mmcv.sparse_flow_from_bytes(flow_bytes)
+
+    # test flow shape is [H, W, 2] and valid shape is [H, W]
+    assert flow_frombytes.shape[:2] == valid_frombytes.shape
+    assert flow_frombytes.shape[2] == 2
+
+    def read_sparse_flow_from_file():
+        flow = cv2.imread(flow_file, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR)
+        flow = flow[:, :, ::-1].astype(np.float32)
+        flow, valid = flow[:, :, :2], flow[:, :, 2]
+        flow = (flow - 2**15) / 64.0
+        return flow, valid
+
+    # read flow from file
+    flow_flowfile, valid_fromfile = read_sparse_flow_from_file()
+
+    assert np.all(flow_frombytes == flow_flowfile)
+    assert np.all(valid_frombytes == valid_fromfile)