Merge pull request #15 from FelixNgFender/main

Examples and restructure python server code

.vscode/settings.json

@@ -1,42 +1,38 @@
 {
-    "notebook.formatOnSave.enabled": true,
-    "notebook.codeActionsOnSave": {
-        "notebook.source.fixAll":"explicit",
-        "notebook.source.organizeImports": "explicit"
+  "notebook.formatOnSave.enabled": true,
+  "notebook.codeActionsOnSave": {
+    "notebook.source.fixAll": "explicit",
+    "notebook.source.organizeImports": "explicit"
+  },
+  "[python]": {
+    "editor.formatOnSave": true,
+    "editor.codeActionsOnSave": {
+      "source.fixAll": "explicit",
+      "source.organizeImports": "explicit"
     },
-    "[python]": {
-        "editor.formatOnSave": true,
-        "editor.codeActionsOnSave": {
-            "source.fixAll": "explicit",
-            "source.organizeImports": "explicit"
-        },
-        "editor.defaultFormatter": "charliermarsh.ruff",
-    },
-    "python.analysis.typeCheckingMode": "off", // TODO: Enable strict type checking
-    "python.analysis.autoImportCompletions": true,
-    "[csharp]": {
-        "editor.formatOnSave": true,
-        "editor.maxTokenizationLineLength": 2500,
-        "editor.inlineSuggest.suppressSuggestions": false
-    },
-    "cSpell.words": [
+    "editor.defaultFormatter": "charliermarsh.ruff"
+  },
+  "python.analysis.typeCheckingMode": "off", // TODO: Enable strict type checking
+  "python.analysis.autoImportCompletions": true,
+  "[csharp]": {
+    "editor.formatOnSave": true,
+    "editor.maxTokenizationLineLength": 2500,
+    "editor.inlineSuggest.suppressSuggestions": false
+  },
+  "cSpell.words": [
         "arange",
-        "arflow",
-        "astype",
-        "Behaviour",
+    "arflow",
+    "astype",
+    "Behaviour",
         "dtype",
-        "flipud",
-        "frombuffer",
-        "gmtime",
+    "flipud",
+    "frombuffer",
+    "gmtime",
         "meshgrid",
-        "ndarray",
-        "Protobuf",
-        "thecakelab",
-        "Xihe"
-    ],
-    "conventionalCommits.scopes": [
-        "server",
-        "client",
-        "others"
-    ],
+    "ndarray",
+    "Protobuf",
+    "thecakelab",
+    "Xihe"
+  ],
+  "conventionalCommits.scopes": ["server", "client", "examples", "others"]
 }

python/README.md

@@ -16,7 +16,7 @@ The ARFlow server can be simply installed via `pip`:
 pip install arflow
 ```
 
-## Example
+## Examples
 
 Next, you may integrate ARFlow with your own research prototype via the Python API:
 
@@ -80,6 +80,8 @@ Save the above code to a file, e.g., `simple_replay.py`, replace `FRAME_DATA_PAT
 python3 simple_replay.py
 ```
 
+For more examples, check out the [examples](https://github.com/cake-lab/ARFlow/tree/main/python/examples) directory.
+
 ## Contributing
 
-We welcome contributions to ARFlow! Please refer to the [CONTRIBUTING.md](https://github.com/cake-lab/ARFlow/blob/main/CONTRIBUTING.md) file for more information.
+We welcome contributions to ARFlow! Please refer to the [`CONTRIBUTING.md`](https://github.com/cake-lab/ARFlow/blob/main/CONTRIBUTING.md) file for more information.

python/arflow/core.py

@@ -1,6 +1,5 @@
 """Data exchanging service."""
 
-import asyncio
 import os
 import pickle
 import time
@@ -20,13 +19,11 @@
 class ARFlowService(service_pb2_grpc.ARFlowService):
     """ARFlow gRPC service."""
 
-    _num_frame: int = 0
-    _loop = asyncio.get_event_loop()
-
     _start_time = time.time_ns()
     _frame_data: List[Dict[str, float | bytes]] = []
 
     def __init__(self) -> None:
+        self.recorder = rr
         super().__init__()
 
     def _save_frame_data(
@@ -51,14 +48,18 @@ def register(
 
         sessions[uid] = request
 
-        rr.init(f"{request.device_name} - ARFlow", spawn=True)
-
+        self.recorder.init(f"{request.device_name} - ARFlow", spawn=True)
         print("Registered a client with UUID: %s" % uid, request)
 
+        # Call the for user extension code.
+        self.on_register(request)
+
         return service_pb2.RegisterResponse(uid=uid)
 
     def data_frame(
-        self, request: service_pb2.DataFrameRequest, context
+        self,
+        request: service_pb2.DataFrameRequest,
+        context,
     ) -> service_pb2.DataFrameResponse:
         """Process an incoming frame."""
 
@@ -72,17 +73,17 @@ def data_frame(
             color_rgb = ARFlowService.decode_rgb_image(session_configs, request.color)
             decoded_data["color_rgb"] = color_rgb
             color_rgb = np.flipud(color_rgb)
-            rr.log("rgb", rr.Image(color_rgb))
+            self.recorder.log("rgb", rr.Image(color_rgb))
 
         if session_configs.camera_depth.enabled:
             depth_img = ARFlowService.decode_depth_image(session_configs, request.depth)
             decoded_data["depth_img"] = depth_img
             depth_img = np.flipud(depth_img)
-            rr.log("depth", rr.DepthImage(depth_img, meter=1.0))
+            self.recorder.log("depth", rr.DepthImage(depth_img, meter=1.0))
 
         if session_configs.camera_transform.enabled:
-            rr.log("world/origin", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN)
-            # rr.log(
+            self.recorder.log("world/origin", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN)
+            # self.logger.log(
             #     "world/xyz",
             #     rr.Arrows3D(
             #         vectors=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
@@ -92,34 +93,42 @@ def data_frame(
 
             transform = ARFlowService.decode_transform(request.transform)
             decoded_data["transform"] = transform
-            rr.log(
+            self.recorder.log(
                 "world/camera",
-                rr.Transform3D(mat3x3=transform[:3, :3], translation=transform[:3, 3]),
+                self.recorder.Transform3D(
+                    mat3x3=transform[:3, :3], translation=transform[:3, 3]
+                ),
             )
 
             k = ARFlowService.decode_intrinsic(session_configs)
-            rr.log("world/camera", rr.Pinhole(image_from_camera=k))
-            rr.log("world/camera", rr.Image(np.flipud(color_rgb)))
+            self.recorder.log("world/camera", rr.Pinhole(image_from_camera=k))
+            self.recorder.log("world/camera", rr.Image(np.flipud(color_rgb)))
 
         if session_configs.camera_point_cloud.enabled:
             pcd, clr = ARFlowService.decode_point_cloud(
                 session_configs, k, color_rgb, depth_img, transform
             )
             decoded_data["point_cloud_pcd"] = pcd
             decoded_data["point_cloud_clr"] = clr
-            rr.log("world/point_cloud", rr.Points3D(pcd, colors=clr))
+            self.recorder.log("world/point_cloud", rr.Points3D(pcd, colors=clr))
 
         # Call the for user extension code.
         self.on_frame_received(decoded_data)
 
         return service_pb2.DataFrameResponse(message="OK")
 
+    def on_register(self, request: service_pb2.RegisterRequest):
+        """Called when a new device is registered. Override this method to process the data."""
+        pass
+
     def on_frame_received(self, frame_data: service_pb2.DataFrameRequest):
         """Called when a frame is received. Override this method to process the data."""
         pass
 
-    def on_program_exit(self, path_to_save: str):
+    def on_program_exit(self, path_to_save: str | None):
         """Save the data and exit."""
+        if path_to_save is None:
+            return
         print("Saving the data...")
         f_name = strftime("%Y_%m_%d_%H_%M_%S", gmtime())
         save_path = os.path.join(path_to_save, f"frames_{f_name}.pkl")
@@ -264,4 +273,4 @@ def decode_point_cloud(
         pcd = np.matmul(transform[:3, :3], pcd.T).T + transform[:3, 3]
         clr = color_rgb.reshape(-1, 3)
 
-        return pcd, clr
+        return pcd, clr

python/examples/README.md

@@ -0,0 +1,51 @@
+# ARFlow Server Examples
+
+The simplest example is [`minimal`](minimal/minimal.py). You may want to start there!
+
+## Setup
+
+If you're using `pip`, you should create and activate a virtual environment before installing any example's dependencies:
+
+```sh
+python3 -m venv .venv
+source .venv/bin/activate
+```
+
+If you're using `poetry` instead, you can just install the dependencies directly, as shown below.
+
+## Installing the example
+
+Each example is packaged as a regular Python package, with a `pyproject.toml` file specifying its required dependencies. To run an example, it must first be installed.
+
+For example, to install dependencies and run the toy `minimal` example (which doesn't need to download any data) run:
+
+```sh
+# Using pip:
+pip install -e python/examples/minimal
+
+# Using poetry:
+cd python/examples/minimal
+poetry install
+```
+
+**Note**: it is import to install example in editable mode, which is done using the `-e` flag (short for `--editable`).
+
+Once installed, the example can be run as a regular Python module:
+
+```shell
+python3 -m minimal
+
+# or, if you're using poetry:
+poetry run minimal
+```
+
+Examples also declare console script, so they can also be run directly:
+
+```shell
+minimal
+```
+
+## Contributions welcome
+Feel free to open a PR to add a new example!
+
+See the [`CONTRIBUTING.md`](https://github.com/cake-lab/ARFlow/blob/main/CONTRIBUTING.md) file  for details on how to contribute.

python/examples/__init__.py

@@ -0,0 +1,3 @@
+"""
+.. include:: ./README.md
+"""

python/examples/depthanythingv2/README.md

@@ -0,0 +1 @@
+# TODO

python/examples/depthanythingv2/depthanythingv2.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Demonstrates the usage of ARFlow with Depth Anything v2."""
+
+from __future__ import annotations
+
+import sys
+from threading import Thread
+from typing import Any, Dict
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import pipeline
+
+import arflow
+
+
+class DepthAnythingV2Service(arflow.ARFlowService):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.pipe = pipeline(
+            "depth-estimation",
+            model="depth-anything/Depth-Anything-V2-base-hf",
+            device=self.device,
+        )
+
+    def on_register(self, request: arflow.RegisterRequest):
+        self.num_frame = 0
+
+    def on_frame_received(self, frame_data: Dict[str, Any]):
+        color_rgb = frame_data["color_rgb"]
+        if self.num_frame % 50 == 0:
+            thread = Thread(target=lambda: (self.run_depth_estimation(color_rgb.copy())) )
+            thread.start()
+
+        self.num_frame = self.num_frame + 1
+
+    def run_depth_estimation(self, color_rgb: np.ndarray):
+        """Run depth estimation on the given image. The pipeline returns a dictionary with two entries.
+        The first one, called predicted_depth, is a tensor with the values being the depth expressed in
+        meters for each pixel. The second one, depth, is a PIL image that visualizes the depth estimation result."""
+
+        image = Image.fromarray(np.flipud(color_rgb))
+
+        predictions = self.pipe(image)
+        self.record_predictions(predictions)
+        return predictions
+
+    def record_predictions(self, predictions: dict):
+        self.recorder.log(
+            "DepthAnythingV2/depth", self.recorder.Image(predictions["depth"])
+        )
+
+
+def main() -> None:
+    # sanity-check since all other example scripts take arguments:
+    assert len(sys.argv) == 1, f"{sys.argv[0]} does not take any arguments"
+    arflow.create_server(DepthAnythingV2Service, port=8500, path_to_save=None)
+
+
+if __name__ == "__main__":
+    main()