|
|
import matplotlib.pyplot as plt |
|
|
import matplotlib.gridspec as gridspec |
|
|
from matplotlib.colors import hsv_to_rgb |
|
|
import cv2 |
|
|
import moviepy.editor as mpy |
|
|
import numpy as np |
|
|
from nowcasting.helpers.gifmaker import save_gif |
|
|
|
|
|
def flow_to_img(flow_dat, max_displacement=None): |
|
|
"""Convert optical flow data to HSV images |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
flow_dat : np.ndarray |
|
|
Shape: (seq_len, 2, H, W) |
|
|
max_displacement : float or None |
|
|
|
|
|
Returns |
|
|
------- |
|
|
rgb_dat : np.ndarray |
|
|
Shape: (seq_len, 3, H, W) |
|
|
""" |
|
|
assert flow_dat.ndim == 4 |
|
|
flow_scale = np.square(flow_dat).sum(axis=1, keepdims=True) |
|
|
flow_x = flow_dat[:, :1, :, :] |
|
|
flow_y = flow_dat[:, 1:, :, :] |
|
|
flow_angle = np.arctan2(flow_y, flow_x) |
|
|
flow_angle[flow_angle < 0] += np.pi * 2 |
|
|
v = np.ones((flow_dat.shape[0], 1, flow_dat.shape[2], flow_dat.shape[3]), |
|
|
dtype=np.float32) |
|
|
if max_displacement is None: |
|
|
flow_scale_max = np.sqrt(flow_scale.max()) |
|
|
else: |
|
|
flow_scale_max = max_displacement |
|
|
h = flow_angle / (2 * np.pi) |
|
|
s = np.sqrt(flow_scale) / flow_scale_max |
|
|
|
|
|
hsv_dat = np.concatenate((h, s, v), axis=1) |
|
|
rgb_dat = hsv_to_rgb(hsv_dat.transpose((0, 2, 3, 1))).transpose((0, 3, 1, 2)) |
|
|
return rgb_dat |
|
|
|
|
|
|
|
|
def _ax_imshow(ax, im, **kwargs): |
|
|
assert im.ndim == 3 or im.ndim == 2 |
|
|
if im.ndim == 2: |
|
|
ax.imshow(im, **kwargs) |
|
|
ax.set_axis_off() |
|
|
else: |
|
|
if im.shape[0] == 1: |
|
|
ax.imshow(im[0, :, :], **kwargs) |
|
|
ax.set_axis_off() |
|
|
elif im.shape[0] == 3: |
|
|
ax.imshow(im.transpose((1, 2, 0)), **kwargs) |
|
|
ax.set_axis_off() |
|
|
else: |
|
|
raise NotImplementedError |
|
|
ax.set_adjustable('box-forced') |
|
|
ax.autoscale(False) |
|
|
|
|
|
|
|
|
def get_color_flow_legend_image(size=50): |
|
|
U, V = np.meshgrid(np.arange(-size, size + 1, dtype=np.float32), |
|
|
np.arange(-size, size + 1, dtype=np.float32)) |
|
|
flow_scale = np.sqrt(U**2 + V**2) |
|
|
flow_angle = np.arctan2(V, U) |
|
|
flow_angle[flow_angle < 0] += np.pi * 2 |
|
|
max_flow_scale = float(size) * np.sqrt(2) |
|
|
h = flow_angle / (2 * np.pi) |
|
|
s = flow_scale / max_flow_scale |
|
|
v = np.ones((size * 2 + 1, size * 2 + 1), |
|
|
dtype=np.float32) |
|
|
hsv_dat = np.concatenate((h.reshape((1, size * 2 + 1, size * 2 + 1)), |
|
|
s.reshape((1, size * 2 + 1, size * 2 + 1)), |
|
|
v.reshape((1, size * 2 + 1, size * 2 + 1))), axis=0) |
|
|
rgb_dat = hsv_to_rgb(hsv_dat.transpose((1, 2, 0))).transpose((2, 0, 1)) |
|
|
a = np.ones((1, size * 2 + 1, size * 2 + 1), dtype=np.float32) |
|
|
rgb_dat[:, flow_scale > max_flow_scale] = 0 |
|
|
a[:, flow_scale > max_flow_scale] = 0 |
|
|
rgba_dat = np.concatenate((rgb_dat, a), axis=0) |
|
|
return rgba_dat |
|
|
|
|
|
|
|
|
def save_hko_gif(im_dat, save_path): |
|
|
"""Save the HKO images to gif |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
im_dat : np.ndarray |
|
|
Shape: (seqlen, H, W) |
|
|
save_path : str |
|
|
Returns |
|
|
------- |
|
|
""" |
|
|
assert im_dat.ndim == 3 |
|
|
save_gif(im_dat, fname=save_path) |
|
|
return |
|
|
|
|
|
|
|
|
def merge_rgba_cv2(front_img, back_img): |
|
|
"""Merge the front image with the background image using the `Painter's algorithm` |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
front_img : np.ndarray |
|
|
back_img : np.ndarray |
|
|
|
|
|
Returns |
|
|
------- |
|
|
result_img : np.ndarray |
|
|
""" |
|
|
assert front_img.shape == back_img.shape |
|
|
if front_img.dtype == np.uint8: |
|
|
front_img = front_img.astype(np.float32) / 255.0 |
|
|
if back_img.dtype == np.uint8: |
|
|
back_img = back_img.astype(np.float32) / 255.0 |
|
|
result_img = np.zeros(front_img.shape, dtype=np.float32) |
|
|
result_img[:, :, 3] = front_img[:, :, 3] + back_img[:, :, 3] * (1 - front_img[:, :, 3]) |
|
|
result_img[:, :, :3] = (front_img[:, :, :3] * front_img[:, :, 3:] + |
|
|
back_img[:, :, :3] * back_img[:, :, 3:] * (1 - front_img[:, :, 3:])) /\ |
|
|
result_img[:, :, 3:] |
|
|
result_img = (result_img * 255.0).astype(np.uint8) |
|
|
return result_img |
|
|
|
|
|
|
|
|
def save_hko_movie(im_dat, datetime_list, mask_dat=None, save_path="hko.mp4", masked=False, |
|
|
fps=5, prediction_start=None): |
|
|
"""Save the HKO images to a video file |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
im_dat : np.ndarray |
|
|
Shape : (seq_len, H, W) |
|
|
datetime_list : list |
|
|
list of datetimes |
|
|
mask_dat : np.ndarray or None |
|
|
Shape : (seq_len, H, W) |
|
|
save_path : str |
|
|
masked : bool |
|
|
whether the mask the inputs when saving the image |
|
|
fps : float |
|
|
the fps of the saved movie |
|
|
prediction_start : int or None |
|
|
The starting point of the prediction |
|
|
""" |
|
|
from nowcasting.config import cfg |
|
|
central_region = cfg.HKO.EVALUATION.CENTRAL_REGION |
|
|
seq_len, height, width = im_dat.shape |
|
|
display_im_dat = [] |
|
|
mask_color = np.array((0, 170, 160, 150), dtype=np.float32) / 255.0 |
|
|
if im_dat.dtype == np.float32: |
|
|
im_dat = (im_dat * 255).astype(np.uint8) |
|
|
for i in range(im_dat.shape[0]): |
|
|
if not masked: |
|
|
color_im_dat = cv2.cvtColor(im_dat[i], cv2.COLOR_GRAY2RGBA) |
|
|
im = color_im_dat |
|
|
else: |
|
|
im = im_dat[i] * mask_dat[i] |
|
|
im = cv2.cvtColor(im, cv2.COLOR_GRAY2RGBA) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if prediction_start is not None and i >= prediction_start: |
|
|
cv2.putText(im, text=datetime_list[i].strftime('%Y/%m/%d %H:%M'), |
|
|
org=(0, 20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, |
|
|
color=(255, 0, 0, 0)) |
|
|
else: |
|
|
cv2.putText(im, text=datetime_list[i].strftime('%Y/%m/%d %H:%M'), |
|
|
org=(0, 20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, |
|
|
color=(255, 255, 255, 0)) |
|
|
cv2.rectangle(im, |
|
|
pt1=(central_region[0], central_region[1]), |
|
|
pt2=(central_region[2], central_region[3]), |
|
|
color=(0, 255, 0, 0)) |
|
|
display_im_dat.append(im) |
|
|
clip = mpy.ImageSequenceClip(display_im_dat, with_mask=False, fps=fps) |
|
|
clip.write_videofile(save_path, audio=False, verbose=False, threads=4) |
|
|
|
