<reponame>adarshkosta/ssl_e2vid
import os
import cv2
import matplotlib
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F


class Visualization:
    """
    Utility class for the visualization and storage of rendered image-like representation
    of multiple elements of the optical flow estimation and image reconstruction pipeline.
    """

    def __init__(self, kwargs, eval_id=-1):
        self.img_idx = 0
        self.px = kwargs["vis"]["px"]
        self.color_scheme = "green_red"  # gray / blue_red / green_red

        if eval_id >= 0:
            self.store_dir = kwargs["trained_model"] + "results/"
            self.store_dir = self.store_dir + "eval_" + str(eval_id) + "/"
            if not os.path.exists(self.store_dir):
                os.makedirs(self.store_dir)
            self.store_file = None

    def update(self, inputs, flow, iwe, brightness):
        """
        Live visualization.
        :param inputs: dataloader dictionary
        :param flow: [batch_size x 2 x H x W] optical flow map
        :param iwe: [batch_size x 1 x H x W] image of warped events
        :param brightness: [batch_size x 1 x H x W] reconstructed image
        """

        inp_events = inputs["inp_cnt"] if "inp_cnt" in inputs.keys() else None
        inp_frames = inputs["inp_frames"] if "inp_frames" in inputs.keys() else None
        height = inp_events.shape[2]
        width = inp_events.shape[3]

        # input events
        inp_events = inp_events.detach()
        inp_events_npy = inp_events.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, -1))
        cv2.namedWindow("Input Events", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Input Events", int(self.px), int(self.px))
        cv2.imshow("Input Events", self.events_to_image(inp_events_npy))

        # input frames
        if inp_frames is not None:
            frame_image = np.zeros((height, 2 * width))
            inp_frames = inp_frames.detach()
            inp_frames_npy = inp_frames.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            frame_image[:height, 0:width] = inp_frames_npy[:, :, 0] / 255.0
            frame_image[:height, width : 2 * width] = inp_frames_npy[:, :, 1] / 255.0
            cv2.namedWindow("Input Frames (Prev/Curr)", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Input Frames (Prev/Curr)", int(2 * self.px), int(self.px))
            cv2.imshow("Input Frames (Prev/Curr)", frame_image)

        # optical flow
        if flow is not None:
            flow = flow.detach()
            flow_npy = flow.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            flow_npy = self.flow_to_image(flow_npy[:, :, 0], flow_npy[:, :, 1])
            flow_npy = cv2.cvtColor(flow_npy, cv2.COLOR_RGB2BGR)
            cv2.namedWindow("Estimated Flow", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Estimated Flow", int(self.px), int(self.px))
            cv2.imshow("Estimated Flow", flow_npy)

        # image of warped events
        if iwe is not None:
            iwe = iwe.detach()
            iwe_npy = iwe.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            iwe_npy = self.events_to_image(iwe_npy)
            cv2.namedWindow("Image of Warped Events", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Image of Warped Events", int(self.px), int(self.px))
            cv2.imshow("Image of Warped Events", iwe_npy)

        # reconstructed brightness
        if brightness is not None:
            brightness = brightness.detach()
            brightness_npy = brightness.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 1))
            intensity_npy = brightness_npy.reshape((height, width, 1))
            intensity_image = self.minmax_norm(intensity_npy)
            cv2.namedWindow("Reconstructed Intensity", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Reconstructed Intensity", int(self.px), int(self.px))
            cv2.imshow("Reconstructed Intensity", intensity_image)

        cv2.waitKey(1)

    def store(self, inputs, flow, iwe, brightness, sequence, ts=None):
        """
        Store rendered images.
        :param inputs: dataloader dictionary
        :param flow: [batch_size x 2 x H x W] optical flow map
        :param iwe: [batch_size x 1 x H x W] image of warped events
        :param brightness: [batch_size x 1 x H x W] reconstructed image
        :param sequence: filename of the event sequence under analysis
        :param ts: timestamp associated with rendered files (default = None)
        """

        inp_events = inputs["inp_cnt"] if "inp_cnt" in inputs.keys() else None
        inp_frames = inputs["inp_frames"] if "inp_frames" in inputs.keys() else None
        height = inp_events.shape[2]
        width = inp_events.shape[3]

        # check if new sequence
        path_to = self.store_dir + sequence + "/"
        if not os.path.exists(path_to):
            os.makedirs(path_to)
            os.makedirs(path_to + "events/")
            os.makedirs(path_to + "flow/")
            os.makedirs(path_to + "frames/")
            os.makedirs(path_to + "iwe/")
            os.makedirs(path_to + "brightness/")
            if self.store_file is not None:
                self.store_file.close()
            self.store_file = open(path_to + "timestamps.txt", "w")
            self.img_idx = 0

        # input events
        event_image = np.zeros((height, width))
        inp_events = inp_events.detach()
        inp_events_npy = inp_events.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, -1))
        event_image = self.events_to_image(inp_events_npy)
        filename = path_to + "events/%09d.png" % self.img_idx
        cv2.imwrite(filename, event_image * 255)

        # input frames
        if inp_frames is not None:
            inp_frames = inp_frames.detach()
            inp_frames_npy = inp_frames.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            filename = path_to + "frames/%09d.png" % self.img_idx
            cv2.imwrite(filename, inp_frames_npy[:, :, 1])

        # optical flow
        if flow is not None:
            flow = flow.detach()
            flow_npy = flow.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            flow_npy = self.flow_to_image(flow_npy[:, :, 0], flow_npy[:, :, 1])
            flow_npy = cv2.cvtColor(flow_npy, cv2.COLOR_RGB2BGR)
            filename = path_to + "flow/%09d.png" % self.img_idx
            cv2.imwrite(filename, flow_npy)

        # image of warped events
        if iwe is not None:
            iwe = iwe.detach()
            iwe_npy = iwe.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            iwe_npy = self.events_to_image(iwe_npy)
            filename = path_to + "iwe/%09d.png" % self.img_idx
            cv2.imwrite(filename, iwe_npy * 255)

        # reconstructed brightness
        if brightness is not None:
            brightness = brightness.detach()
            brightness_npy = brightness.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 1))
            intensity_npy = brightness_npy.reshape((height, width, 1))
            intensity_image = self.minmax_norm(intensity_npy)
            filename = path_to + "brightness/%09d.png" % self.img_idx
            cv2.imwrite(filename, intensity_image * 255)

        # store timestamps
        if ts is not None:
            self.store_file.write(str(ts) + "\n")
            self.store_file.flush()

        self.img_idx += 1
        cv2.waitKey(1)

    @staticmethod
    def flow_to_image(flow_x, flow_y):
        """
        Use the optical flow color scheme from the supplementary materials of the paper 'Back to Event
        Basics: Self-Supervised Image Reconstruction for Event Cameras via Photometric Constancy',
        Paredes-Valles et al., CVPR'21.
        :param flow_x: [H x W x 1] horizontal optical flow component
        :param flow_y: [H x W x 1] vertical optical flow component
        :return flow_rgb: [H x W x 3] color-encoded optical flow
        """
        flows = np.stack((flow_x, flow_y), axis=2)
        mag = np.linalg.norm(flows, axis=2)
        min_mag = np.min(mag)
        mag_range = np.max(mag) - min_mag

        ang = np.arctan2(flow_y, flow_x) + np.pi
        ang *= 1.0 / np.pi / 2.0

        hsv = np.zeros([flow_x.shape[0], flow_x.shape[1], 3])
        hsv[:, :, 0] = ang
        hsv[:, :, 1] = 1.0
        hsv[:, :, 2] = mag - min_mag
        if mag_range != 0.0:
            hsv[:, :, 2] /= mag_range

        flow_rgb = matplotlib.colors.hsv_to_rgb(hsv)
        return (255 * flow_rgb).astype(np.uint8)

    @staticmethod
    def minmax_norm(x):
        """
        Robust min-max normalization.
        :param x: [H x W x 1]
        :return x: [H x W x 1] normalized x
        """
        den = np.percentile(x, 99) - np.percentile(x, 1)
        if den != 0:
            x = (x - np.percentile(x, 1)) / den
        return np.clip(x, 0, 1)

    @staticmethod
    def events_to_image(inp_events, color_scheme="green_red"):
        """
        Visualize the input events.
        :param inp_events: [batch_size x 2 x H x W] per-pixel and per-polarity event count
        :param color_scheme: green_red/gray
        :return event_image: [H x W x 3] color-coded event image
        """
        pos = inp_events[:, :, 0]
        neg = inp_events[:, :, 1]
        pos_max = np.percentile(pos, 99)
        pos_min = np.percentile(pos, 1)
        neg_max = np.percentile(neg, 99)
        neg_min = np.percentile(neg, 1)
        max = pos_max if pos_max > neg_max else neg_max

        if pos_min != max:
            pos = (pos - pos_min) / (max - pos_min)
        if neg_min != max:
            neg = (neg - neg_min) / (max - neg_min)

        pos = np.clip(pos, 0, 1)
        neg = np.clip(neg, 0, 1)

        event_image = np.ones((inp_events.shape[0], inp_events.shape[1]))
        if color_scheme == "gray":
            event_image *= 0.5
            pos *= 0.5
            neg *= -0.5
            event_image += pos + neg

        elif color_scheme == "green_red":
            event_image = np.repeat(event_image[:, :, np.newaxis], 3, axis=2)
            event_image *= 0
            mask_pos = pos > 0
            mask_neg = neg > 0
            mask_not_pos = pos == 0
            mask_not_neg = neg == 0

            event_image[:, :, 0][mask_pos] = 0
            event_image[:, :, 1][mask_pos] = pos[mask_pos]
            event_image[:, :, 2][mask_pos * mask_not_neg] = 0
            event_image[:, :, 2][mask_neg] = neg[mask_neg]
            event_image[:, :, 0][mask_neg] = 0
            event_image[:, :, 1][mask_neg * mask_not_pos] = 0

        return event_image

import os
import cv2
import matplotlib
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F


class Visualization:
    """
    Utility class for the visualization and storage of rendered image-like representation
    of multiple elements of the optical flow estimation and image reconstruction pipeline.
    """

    def __init__(self, kwargs, eval_id=-1):
        self.img_idx = 0
        self.px = kwargs["vis"]["px"]
        self.color_scheme = "green_red"  # gray / blue_red / green_red

        if eval_id >= 0:
            self.store_dir = kwargs["trained_model"] + "results/"
            self.store_dir = self.store_dir + "eval_" + str(eval_id) + "/"
            if not os.path.exists(self.store_dir):
                os.makedirs(self.store_dir)
            self.store_file = None

    def update(self, inputs, flow, iwe, brightness):
        """
        Live visualization.
        :param inputs: dataloader dictionary
        :param flow: [batch_size x 2 x H x W] optical flow map
        :param iwe: [batch_size x 1 x H x W] image of warped events
        :param brightness: [batch_size x 1 x H x W] reconstructed image
        """

        inp_events = inputs["inp_cnt"] if "inp_cnt" in inputs.keys() else None
        inp_frames = inputs["inp_frames"] if "inp_frames" in inputs.keys() else None
        height = inp_events.shape[2]
        width = inp_events.shape[3]

        # input events
        inp_events = inp_events.detach()
        inp_events_npy = inp_events.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, -1))
        cv2.namedWindow("Input Events", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Input Events", int(self.px), int(self.px))
        cv2.imshow("Input Events", self.events_to_image(inp_events_npy))

        # input frames
        if inp_frames is not None:
            frame_image = np.zeros((height, 2 * width))
            inp_frames = inp_frames.detach()
            inp_frames_npy = inp_frames.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            frame_image[:height, 0:width] = inp_frames_npy[:, :, 0] / 255.0
            frame_image[:height, width : 2 * width] = inp_frames_npy[:, :, 1] / 255.0
            cv2.namedWindow("Input Frames (Prev/Curr)", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Input Frames (Prev/Curr)", int(2 * self.px), int(self.px))
            cv2.imshow("Input Frames (Prev/Curr)", frame_image)

        # optical flow
        if flow is not None:
            flow = flow.detach()
            flow_npy = flow.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            flow_npy = self.flow_to_image(flow_npy[:, :, 0], flow_npy[:, :, 1])
            flow_npy = cv2.cvtColor(flow_npy, cv2.COLOR_RGB2BGR)
            cv2.namedWindow("Estimated Flow", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Estimated Flow", int(self.px), int(self.px))
            cv2.imshow("Estimated Flow", flow_npy)

        # image of warped events
        if iwe is not None:
            iwe = iwe.detach()
            iwe_npy = iwe.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            iwe_npy = self.events_to_image(iwe_npy)
            cv2.namedWindow("Image of Warped Events", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Image of Warped Events", int(self.px), int(self.px))
            cv2.imshow("Image of Warped Events", iwe_npy)

        # reconstructed brightness
        if brightness is not None:
            brightness = brightness.detach()
            brightness_npy = brightness.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 1))
            intensity_npy = brightness_npy.reshape((height, width, 1))
            intensity_image = self.minmax_norm(intensity_npy)
            cv2.namedWindow("Reconstructed Intensity", cv2.WINDOW_NORMAL)
            cv2.resizeWindow("Reconstructed Intensity", int(self.px), int(self.px))
            cv2.imshow("Reconstructed Intensity", intensity_image)

        cv2.waitKey(1)

    def store(self, inputs, flow, iwe, brightness, sequence, ts=None):
        """
        Store rendered images.
        :param inputs: dataloader dictionary
        :param flow: [batch_size x 2 x H x W] optical flow map
        :param iwe: [batch_size x 1 x H x W] image of warped events
        :param brightness: [batch_size x 1 x H x W] reconstructed image
        :param sequence: filename of the event sequence under analysis
        :param ts: timestamp associated with rendered files (default = None)
        """

        inp_events = inputs["inp_cnt"] if "inp_cnt" in inputs.keys() else None
        inp_frames = inputs["inp_frames"] if "inp_frames" in inputs.keys() else None
        height = inp_events.shape[2]
        width = inp_events.shape[3]

        # check if new sequence
        path_to = self.store_dir + sequence + "/"
        if not os.path.exists(path_to):
            os.makedirs(path_to)
            os.makedirs(path_to + "events/")
            os.makedirs(path_to + "flow/")
            os.makedirs(path_to + "frames/")
            os.makedirs(path_to + "iwe/")
            os.makedirs(path_to + "brightness/")
            if self.store_file is not None:
                self.store_file.close()
            self.store_file = open(path_to + "timestamps.txt", "w")
            self.img_idx = 0

        # input events
        event_image = np.zeros((height, width))
        inp_events = inp_events.detach()
        inp_events_npy = inp_events.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, -1))
        event_image = self.events_to_image(inp_events_npy)
        filename = path_to + "events/%09d.png" % self.img_idx
        cv2.imwrite(filename, event_image * 255)

        # input frames
        if inp_frames is not None:
            inp_frames = inp_frames.detach()
            inp_frames_npy = inp_frames.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            filename = path_to + "frames/%09d.png" % self.img_idx
            cv2.imwrite(filename, inp_frames_npy[:, :, 1])

        # optical flow
        if flow is not None:
            flow = flow.detach()
            flow_npy = flow.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            flow_npy = self.flow_to_image(flow_npy[:, :, 0], flow_npy[:, :, 1])
            flow_npy = cv2.cvtColor(flow_npy, cv2.COLOR_RGB2BGR)
            filename = path_to + "flow/%09d.png" % self.img_idx
            cv2.imwrite(filename, flow_npy)

        # image of warped events
        if iwe is not None:
            iwe = iwe.detach()
            iwe_npy = iwe.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 2))
            iwe_npy = self.events_to_image(iwe_npy)
            filename = path_to + "iwe/%09d.png" % self.img_idx
            cv2.imwrite(filename, iwe_npy * 255)

        # reconstructed brightness
        if brightness is not None:
            brightness = brightness.detach()
            brightness_npy = brightness.cpu().numpy().transpose(0, 2, 3, 1).reshape((height, width, 1))
            intensity_npy = brightness_npy.reshape((height, width, 1))
            intensity_image = self.minmax_norm(intensity_npy)
            filename = path_to + "brightness/%09d.png" % self.img_idx
            cv2.imwrite(filename, intensity_image * 255)

        # store timestamps
        if ts is not None:
            self.store_file.write(str(ts) + "\n")
            self.store_file.flush()

        self.img_idx += 1
        cv2.waitKey(1)

    @staticmethod
    def flow_to_image(flow_x, flow_y):
        """
        Use the optical flow color scheme from the supplementary materials of the paper 'Back to Event
        Basics: Self-Supervised Image Reconstruction for Event Cameras via Photometric Constancy',
        Paredes-Valles et al., CVPR'21.
        :param flow_x: [H x W x 1] horizontal optical flow component
        :param flow_y: [H x W x 1] vertical optical flow component
        :return flow_rgb: [H x W x 3] color-encoded optical flow
        """
        flows = np.stack((flow_x, flow_y), axis=2)
        mag = np.linalg.norm(flows, axis=2)
        min_mag = np.min(mag)
        mag_range = np.max(mag) - min_mag

        ang = np.arctan2(flow_y, flow_x) + np.pi
        ang *= 1.0 / np.pi / 2.0

        hsv = np.zeros([flow_x.shape[0], flow_x.shape[1], 3])
        hsv[:, :, 0] = ang
        hsv[:, :, 1] = 1.0
        hsv[:, :, 2] = mag - min_mag
        if mag_range != 0.0:
            hsv[:, :, 2] /= mag_range

        flow_rgb = matplotlib.colors.hsv_to_rgb(hsv)
        return (255 * flow_rgb).astype(np.uint8)

    @staticmethod
    def minmax_norm(x):
        """
        Robust min-max normalization.
        :param x: [H x W x 1]
        :return x: [H x W x 1] normalized x
        """
        den = np.percentile(x, 99) - np.percentile(x, 1)
        if den != 0:
            x = (x - np.percentile(x, 1)) / den
        return np.clip(x, 0, 1)

    @staticmethod
    def events_to_image(inp_events, color_scheme="green_red"):
        """
        Visualize the input events.
        :param inp_events: [batch_size x 2 x H x W] per-pixel and per-polarity event count
        :param color_scheme: green_red/gray
        :return event_image: [H x W x 3] color-coded event image
        """
        pos = inp_events[:, :, 0]
        neg = inp_events[:, :, 1]
        pos_max = np.percentile(pos, 99)
        pos_min = np.percentile(pos, 1)
        neg_max = np.percentile(neg, 99)
        neg_min = np.percentile(neg, 1)
        max = pos_max if pos_max > neg_max else neg_max

        if pos_min != max:
            pos = (pos - pos_min) / (max - pos_min)
        if neg_min != max:
            neg = (neg - neg_min) / (max - neg_min)

        pos = np.clip(pos, 0, 1)
        neg = np.clip(neg, 0, 1)

        event_image = np.ones((inp_events.shape[0], inp_events.shape[1]))
        if color_scheme == "gray":
            event_image *= 0.5
            pos *= 0.5
            neg *= -0.5
            event_image += pos + neg

        elif color_scheme == "green_red":
            event_image = np.repeat(event_image[:, :, np.newaxis], 3, axis=2)
            event_image *= 0
            mask_pos = pos > 0
            mask_neg = neg > 0
            mask_not_pos = pos == 0
            mask_not_neg = neg == 0

            event_image[:, :, 0][mask_pos] = 0
            event_image[:, :, 1][mask_pos] = pos[mask_pos]
            event_image[:, :, 2][mask_pos * mask_not_neg] = 0
            event_image[:, :, 2][mask_neg] = neg[mask_neg]
            event_image[:, :, 0][mask_neg] = 0
            event_image[:, :, 1][mask_neg * mask_not_pos] = 0

        return event_image
Utility class for the visualization and storage of rendered image-like representation     of multiple elements of the optical flow estimation and image reconstruction pipeline. # gray / blue_red / green_red Live visualization.         :param inputs: dataloader dictionary         :param flow: [batch_size x 2 x H x W] optical flow map         :param iwe: [batch_size x 1 x H x W] image of warped events         :param brightness: [batch_size x 1 x H x W] reconstructed image # input events # input frames # optical flow # image of warped events # reconstructed brightness Store rendered images.         :param inputs: dataloader dictionary         :param flow: [batch_size x 2 x H x W] optical flow map         :param iwe: [batch_size x 1 x H x W] image of warped events         :param brightness: [batch_size x 1 x H x W] reconstructed image         :param sequence: filename of the event sequence under analysis         :param ts: timestamp associated with rendered files (default = None) # check if new sequence # input events # input frames # optical flow # image of warped events # reconstructed brightness # store timestamps Use the optical flow color scheme from the supplementary materials of the paper 'Back to Event         Basics: Self-Supervised Image Reconstruction for Event Cameras via Photometric Constancy',         Paredes-Valles et al., CVPR'21.         :param flow_x: [H x W x 1] horizontal optical flow component         :param flow_y: [H x W x 1] vertical optical flow component         :return flow_rgb: [H x W x 3] color-encoded optical flow Robust min-max normalization.         :param x: [H x W x 1]         :return x: [H x W x 1] normalized x Visualize the input events.         :param inp_events: [batch_size x 2 x H x W] per-pixel and per-polarity event count         :param color_scheme: green_red/gray         :return event_image: [H x W x 3] color-coded event image