Merge pull request #23 from eth-ait/dev

Pull request 1.8.0 release

.github/workflows/tests.yml

@@ -0,0 +1,39 @@
+name: tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.7", "3.10"]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'pip'
+          cache-dependency-path: setup.py
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
+          pip install -e .
+      - name: Setup display
+        if: runner.os == 'Linux'
+        run: |
+          export DISPLAY=:0.0
+          Xvfb :0 -screen 0 640x480x24 &
+      - name: Run tests
+        run: |
+          cd tests
+          pytest

README.md

@@ -1,4 +1,5 @@
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Tests](https://github.com/eth-ait/aitviewer/actions/workflows/tests.yml/badge.svg)](https://github.com/eth-ait/aitviewer/actions/workflows/tests.yml)
 
 # [![aitviewer](assets/aitviewer_logo.svg)](https://github.com/eth-ait/aitviewer)
 A set of tools to visualize and interact with sequences of 3D data with cross-platform support on Windows, Linux, and macOS. See the official page at [https://eth-ait.github.io/aitviewer](https://eth-ait.github.io/aitviewer/) for all the details.

aitviewer/__init__.py

@@ -1 +1 @@
-__version__ = "1.7.1"
+__version__ = "1.8.0"

aitviewer/headless.py

@@ -15,7 +15,9 @@
 along with this program.  If not, see <https://www.gnu.org/licenses/>.
 """
 import os
+from typing import Dict, Tuple
 
+import numpy as np
 from PIL.Image import Image
 
 from aitviewer.viewer import Viewer
@@ -74,15 +76,23 @@ def save_depth(self, file_path):
             os.makedirs(dir, exist_ok=True)
         self.get_depth().save(file_path)
 
-    def save_mask(self, file_path):
+    def save_mask(self, file_path, color_map: Dict[int, Tuple[int, int, int]] = None, id_map: Dict[int, int] = None):
         """
         Render and save a color mask as a 'RGB' PIL image.
+        Each object in the mask has a uniform color computed from the Node UID (can be accessed from a node with 'node.uid').
+
         :param file_path: the path where the image is saved.
+        :param color_map:
+            if not None specifies the color to use for a given Node UID as a tuple (R, G, B) of integer values from 0 to 255.
+            If None the color is computed as an hash of the Node UID instead.
+        :param id_map:
+            if not None the UIDs in the mask are mapped using this dictionary from Node UID to the specified ID.
+            This mapping is applied before the color map (or before hashing if the color map is None).
         """
         dir = os.path.dirname(file_path)
         if dir:
             os.makedirs(dir, exist_ok=True)
-        self.get_mask().save(file_path)
+        self.get_mask(color_map, id_map).save(file_path)
 
     def _render_frame(self):
         self._init_scene()
@@ -114,10 +124,30 @@ def get_depth(self) -> Image:
         self._render_frame()
         return self.get_current_depth_image()
 
-    def get_mask(self) -> Image:
+    def get_mask_ids(self, id_map: Dict[int, int] = None) -> np.ndarray:
+        """
+        Return a mask as a numpy array of shape (height, width) and type np.uint32.
+        Each element in the array is the UID of the node covering that pixel (can be accessed from a node with 'node.uid')
+        or zero if not covered.
+
+        :param id_map:
+            if not None the UIDs in the mask are mapped using this dictionary to the specified ID.
+            The final mask only contains the IDs specified in this mapping and zeros everywhere else.
+        """
+        self._render_frame()
+        return self.get_current_mask_ids(id_map)
+
+    def get_mask(self, color_map: Dict[int, Tuple[int, int, int]] = None, id_map: Dict[int, int] = None) -> Image:
         """
-        Render and return a color mask as a 'RGB' PIL image. Each object in the mask
-        has a uniform color computed as an hash of the Node uid.
+        Render and return a color mask as a 'RGB' PIL image.
+        Each object in the mask has a uniform color computed from the Node UID (can be accessed from a node with 'node.uid').
+
+        :param color_map:
+            if not None specifies the color to use for a given Node UID as a tuple (R, G, B) of integer values from 0 to 255.
+            If None the color is computed as an hash of the Node UID instead.
+        :param id_map:
+            if not None the UIDs in the mask are mapped using this dictionary from Node UID to the specified ID.
+            This mapping is applied before the color map (or before hashing if the color map is None).
         """
         self._render_frame()
-        return self.get_current_mask_image()
+        return self.get_current_mask_image(color_map, id_map)

aitviewer/remote/viewer.py

@@ -1,7 +1,9 @@
 import asyncio
 import pickle
+import queue
 import subprocess
 import threading
+from typing import Callable
 
 import websockets
 
@@ -90,9 +92,6 @@ def _entry(self, url):
     async def _async_entry(self, url):
         # Async entry point of the client thread.
 
-        # Create an async queue for communicating with the main thread.
-        self.queue = asyncio.Queue()
-
         # Attempt to connect until 'self.timeout' seconds passed.
         start_time = self.loop.time()
         try:
@@ -113,29 +112,72 @@ async def _async_entry(self, url):
         if not self.connected:
             return
 
+        # Create a queue for incoming messages to the main thread.
+        self.recv_queue = queue.Queue()
+
         # Message loop.
         try:
-            while True:
-                data = await self.queue.get()
-                if data is None:
-                    await self.websocket.close()
-                    break
-                await self.websocket.send(data)
+            # This loop is exited whenever the connection is dropped
+            # which causes and exception to be raised.
+            async for message in self.websocket:
+                data = pickle.loads(message)
+                # Equeue data for the main thread to process.
+                self.recv_queue.put_nowait(data)
         except Exception as e:
             print(f"Message loop exception: {e}")
 
         # Mark the connection as closed.
         self.connected = False
 
+    def get_message(self, block=True):
+        """
+        Returns the next message received by the remote viewer.
+
+        :param block: if True this function blocks until a message is received, otherwise it returns immediately.
+
+        :return: if block is True returns the next message or None if the connection has been closed.
+                 if block is False returns the next message or None if there are no messages.
+        """
+        if self.connected:
+            if block:
+                while self.connected:
+                    try:
+                        return self.recv_queue.get(timeout=0.1)
+                    except queue.Empty:
+                        pass
+            else:
+                if not self.recv_queue.empty():
+                    return self.recv_queue.get_nowait()
+
+        return None
+
+    def process_messages(self, handler: Callable[["RemoteViewer", object], None], block=True):
+        """
+        Processes messages in a loop calling 'handler' for each message.
+
+        :param block: if True this function blocks until the connection is closed, otherwise it returns
+            after all messages received so far have been processed.
+
+        :return: if block is True always returns False when the connection has been closed.
+                 if block is False returns True if the connection is still open or False if the connection
+                 has been closed.
+        """
+        while True:
+            msg = self.get_message(block)
+            if msg is None:
+                if block:
+                    return False
+                else:
+                    return self.connected
+            handler(self, msg)
+
     async def _async_send(self, data):
-        # Append to the mesage queue.
-        await self.queue.put(data)
+        await self.websocket.send(data)
 
     def send(self, data):
         try:
             if self.connected:
-                # Append a message to the client thread queue by adding a send coroutine to the
-                # thread's loop and wait for it to complete.
+                # Send a message by adding a send coroutine to the thread's loop and wait for it to complete.
                 asyncio.run_coroutine_threadsafe(self._async_send(data), self.loop).result()
         except Exception as e:
             print(f"Send exception: {e}")
@@ -166,11 +208,13 @@ def previous_frame(self):
         """Set the current active frame of the remote viewer to the previous frame"""
         self.send_message(Message.PREVIOUS_FRAME)
 
+    async def _async_close(self):
+        await self.websocket.close()
+
     def close_connection(self):
         """Close the connection with the remote viewer."""
         if self.connected:
-            # Send a special None message to signal the client thread to close the connection.
-            self.send(None)
+            asyncio.run_coroutine_threadsafe(self._async_close(), self.loop).result()
 
             # Wait for the client thread to exit.
             self.thread.join()

aitviewer/renderables/bounding_boxes.py

@@ -51,10 +51,19 @@ def __init__(self, vertices, thickness=0.005, color=(0.0, 0.0, 1.0, 1.0), **kwar
             mode="lines",
             r_base=thickness,
             color=self.color,
+            cast_shadow=False,
         )
-        self.spheres = Spheres(positions=self.vertices, radius=thickness, color=self.color)
+        self.spheres = Spheres(positions=self.vertices, radius=thickness, color=self.color, cast_shadow=False)
         self._add_nodes(self.lines, self.spheres, show_in_hierarchy=False)
 
+    @property
+    def bounds(self):
+        return self.get_bounds(self.vertices)
+
+    @property
+    def current_bounds(self):
+        return self.get_bounds(self.vertices[self.current_frame_id])
+
     @staticmethod
     def from_min_max_diagonal(v_min, v_max, **kwargs):
         """

aitviewer/renderables/lines.py

@@ -234,11 +234,12 @@ def __init__(
         :param r_base: Thickness of the line.
         :param r_tip: If set, the thickness of the line will taper from r_base to r_tip. If set to 0.0 it will create
           a proper cone.
-        :param color: Color of the line (4-tuple).
-        :param mode: 'lines' or 'line_strip' -> ModernGL drawing mode - LINE_STRIP oder LINES
+        :param color: Color of the line (4-tuple) or array of color (N_LINES, 4), one for each line.
+        :param mode: 'lines' or 'line_strip'.
+            'lines': a line is drawn from point 0 to 1, from 2 to 3, and so on, number of lines is L / 2.
+            'line_strip': a line is drawn between all adjacent points, 0 to 1, 1 to 2 and so on, number of lines is L - 1.
         :param cast_shadow: If True the mesh casts a shadow on other objects.
         """
-        assert len(color) == 4
         if len(lines.shape) == 2:
             lines = lines[np.newaxis]
         assert len(lines.shape) == 3
@@ -254,7 +255,16 @@ def __init__(
         self.vertices, self.faces = self.get_mesh()
         self.n_lines = self.lines.shape[1] // 2 if mode == "lines" else self.lines.shape[1] - 1
 
-        kwargs["material"] = kwargs.get("material", Material(color=color, ambient=0.2))
+        # Define a default material in case there is None.
+        if isinstance(color, tuple) or len(color.shape) == 1:
+            kwargs["material"] = kwargs.get("material", Material(color=color, ambient=0.2))
+            self.line_colors = kwargs["material"].color
+        else:
+            assert (
+                color.shape[1] == 4 and color.shape[0] == self.n_lines
+            ), "Color must be a tuple of 4 values or a numpy array of shape (N_LINES, 4)"
+            self.line_colors = color
+
         super(Lines, self).__init__(n_frames=self.lines.shape[0], **kwargs)
 
         self._need_upload = True
@@ -304,6 +314,27 @@ def current_lines(self, lines):
         self._lines[idx] = lines
         self.redraw()
 
+    @Node.color.setter
+    def color(self, color):
+        self.material.color = color
+        self.line_colors = color
+        self.redraw()
+
+    @property
+    def line_colors(self):
+        if len(self._line_colors.shape) == 1:
+            t = np.tile(np.array(self._line_colors), (self.n_lines, 1))
+            return t
+        else:
+            return self._line_colors
+
+    @line_colors.setter
+    def line_colors(self, color):
+        if isinstance(color, tuple):
+            color = np.array(color)
+        self._line_colors = color
+        self.redraw()
+
     def on_frame_update(self):
         self.redraw()
 
@@ -323,11 +354,13 @@ def make_renderable(self, ctx: moderngl.Context):
         self.vbo_indices = ctx.buffer(self.faces.astype("i4").tobytes())
         self.vbo_instance_base = ctx.buffer(reserve=self.n_lines * 12)
         self.vbo_instance_tip = ctx.buffer(reserve=self.n_lines * 12)
+        self.vbo_instance_color = ctx.buffer(reserve=self.n_lines * 16)
 
         self.vao = VAO()
         self.vao.buffer(self.vbo_vertices, "3f4", "in_position")
         self.vao.buffer(self.vbo_instance_base, "3f4/i", "instance_base")
         self.vao.buffer(self.vbo_instance_tip, "3f4/i", "instance_tip")
+        self.vao.buffer(self.vbo_instance_color, "4f4/i", "instance_color")
         self.vao.index_buffer(self.vbo_indices)
 
     def _upload_buffers(self):
@@ -346,16 +379,23 @@ def _upload_buffers(self):
         self.vbo_instance_base.write(v0s.astype("f4").tobytes())
         self.vbo_instance_tip.write(v1s.astype("f4").tobytes())
 
+        if len(self._line_colors.shape) > 1:
+            self.vbo_instance_color.write(self._line_colors.astype("f4").tobytes())
+
     def render(self, camera, **kwargs):
         self._upload_buffers()
 
         prog = self.prog
         prog["r_base"] = self.r_base
         prog["r_tip"] = self.r_tip
-        prog["use_uniform_color"] = True
-        prog["uniform_color"] = tuple(self.color)
+        if len(self._line_colors.shape) == 1:
+            prog["use_uniform_color"] = True
+            prog["uniform_color"] = tuple(self.color)
+        else:
+            prog["use_uniform_color"] = False
         prog["draw_edges"].value = 1.0 if self.draw_edges else 0.0
         prog["win_size"].value = kwargs["window_size"]
+        prog["clip_control"].value = (0, 0, 0)
 
         self.set_camera_matrices(prog, camera, **kwargs)
         set_lights_in_program(