NetCDF to MP4 demo

netcdf_to_mp4_demo.py

import numpy as np
import xarray as xr
import cv2
import imageio


def to_rgb(data, vmin=None, vmax=None, size=None, pmin=2, pmax=98, mask=None):
    if isinstance(data, list):
        n_channels = len(data)
    elif isinstance(data, xr.DataArray):
        n_channels = 1
        data = [data]
    else:
        raise ValueError()

    values = [np.asarray(d) for d in data]
    shape = data[0].shape + (n_channels,)

    if vmin is not None:
        if isinstance(vmin, (int, float)):
            vmin = [vmin] * n_channels
    if vmax is not None:
        if isinstance(vmax, (int, float)):
            vmax = [vmax] * n_channels

    im = np.empty(shape)

    for i in range(n_channels):
        channel = values[i]
        if vmin is not None:
            minval = vmin[i]
        else:
            minval = np.percentile(channel, pmin)
        if vmax is not None:
            maxval = vmax[i]
        else:
            maxval = np.percentile(channel, pmax)
        if maxval > minval:
            channel = (channel - minval) / (maxval - minval) * 255

        im[:, :, i] = channel
    im = np.clip(im, 0, 255).astype(np.uint8)
    if n_channels == 1:
        colored = cv2.cvtColor(im[:, :, 0], cv2.COLOR_GRAY2BGR)
    else:
        colored = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)

    if mask is not None:
        colored[~mask] = 0

    if size is not None:
        if size[0] is None:
            size = (int(colored.shape[0] * size[1] / colored.shape[1]), size[1])
        elif size[1] is None:
            size = (size[0], int(colored.shape[1] * size[0] / colored.shape[0]))
        colored = cv2.resize(colored, (size[1], size[0]))

    return cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)


filenames = ['2014_01_01_2m_temperature.nc',
             '2014_01_01_100m_u_component_of_wind.nc',
             '2014_01_01_100m_v_component_of_wind.nc']

da = xr.open_mfdataset(filenames, concat_dim='time').to_array()
with imageio.get_writer('result.mp4', codec='libx264') as writer:
    for t in da.time.values:
        data = []
        for name in ['t2m', 'u100', 'v100']:
            data.append(da.sel(time=t, variable=name))
        frame = to_rgb(data)
        writer.append_data(frame)