From d4406e817a31533fe8cff846410a6c0f2eb7534d Mon Sep 17 00:00:00 2001
From: shiba24 <shiba.shintaro@gmail.com>
Date: Fri, 30 Jun 2023 15:10:10 +0900
Subject: [PATCH] Visualize events v1

---
 src/evlib/vis/view2d.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/evlib/vis/view2d.py b/src/evlib/vis/view2d.py
index 82d0885..a7e5986 100644
--- a/src/evlib/vis/view2d.py
+++ b/src/evlib/vis/view2d.py
@@ -75,3 +75,28 @@ def optical_flow(
     else:
         wheel = None
     return image, wheel
+
+
+def events(events: np.ndarray, image_shape: tuple) -> Image.Image:
+    """Visualize events with polarity color.
+
+    Args:
+        events (np.ndarray): _description_
+        image_shape (tuple): _description_
+
+    Returns:
+        _type_: _description_
+    """
+    background_color = 255
+    image = (
+        np.ones((image_shape[0], image_shape[1], 3), dtype=np.uint8)
+        * background_color
+    )  # RGBA channel
+
+    events[:, 0] = np.clip(events[:, 0], 0, image_shape[0] - 1)
+    events[:, 1] = np.clip(events[:, 1], 0, image_shape[1] - 1)
+    colors = np.array([(255, 0, 0) if e[3] == 1 else (0, 0, 255) for e in events])
+    image[events[:, 0].astype(np.int32), events[:, 1].astype(np.int32), :] = colors
+
+    image = Image.fromarray(image)
+    return image