Feat/point cloud

demos/PointCloudLabeling/point_cloud_loader.py

@@ -0,0 +1,43 @@
+import random
+import cv2
+
+import json
+from simpub.core.net_manager import init_net_manager
+from simpub.xr_device.xr_device import XRDevice
+
+def generate_point_cloud_data(rgb_image, depth_image):
+    positions = []
+    colors = []
+    height, width, _ = rgb_image.shape
+    for y in range(height):
+        for x in range(width):
+            rgb_pixel = rgb_image[y, x]
+            depth_pixel = depth_image[y, x]
+            if random.random() > 0.1:
+                continue
+            positions.extend([x / width, y / height, depth_pixel[0] / 10])
+            colors.extend([int(rgb_pixel[0]) / 255, int(rgb_pixel[1]) / 255, int(rgb_pixel[2]) / 255, 1])
+    return json.dumps({"positions": positions, "colors": colors})
+
+
+def send_point_cloud(self, point_data) -> None:
+    self.request_socket.send_string(json.dumps(self.point_cloud_data))
+
+
+rgb_path = '/home/xinkai/project/point-cloud-labeling/rgb.png'
+depth_path = '/home/xinkai/project/point-cloud-labeling/depth.png'
+
+rgb_image = cv2.imread(rgb_path)
+depth_image = cv2.imread(depth_path)
+
+# simpublisher = RGBDPublisher(rgb_image, depth_image)
+net_manager = init_net_manager("127.0.0.1")
+unity_editor = XRDevice()
+# print(generate_point_cloud_data(rgb_image, depth_image))
+while unity_editor.connected is False:
+    pass
+unity_editor.request("LoadPointCloud", generate_point_cloud_data(rgb_image, depth_image))
+net_manager.join()
+
+# data = simpublisher.generate_point_cloud_data(rgb_image, depth_image)
+# simpublisher.join()

demos/SimulationFramework/pick_and_place.py

@@ -3,7 +3,7 @@
 from alr_sim.sims.SimFactory import SimRepository
 from alr_sim.sims.universal_sim.PrimitiveObjects import Box
 from simpub.sim.sf_publisher import SFPublisher
-from simpub.xr_device.meta_quest3 import MetaQuest3
+from simpub.xr_device.xr_device import XRDevice
 
 
 if __name__ == "__main__":
@@ -72,7 +72,7 @@
     publisher = SFPublisher(
         scene, host, no_tracked_objects=["table_plane", "table0"]
     )
-    device = MetaQuest3()
+    device = XRDevice()
 
     duration = 2
 

simpub/core/net_manager.py

@@ -8,7 +8,7 @@
 import struct
 from time import sleep
 import json
-
+import uuid
 from .log import logger
 
 IPAddress = NewType("IPAddress", str)
@@ -117,14 +117,15 @@ def broadcast_loop(self):
         # set up udp socket
         _socket = socket.socket(AF_INET, SOCK_DGRAM)
         _socket.setsockopt(SOL_SOCKET, SO_BROADCAST, 1)
+        _id = str(uuid.uuid4())
         # calculate broadcast ip
         local_info = self.local_info
         ip_bin = struct.unpack('!I', socket.inet_aton(local_info["ip"]))[0]
         netmask_bin = struct.unpack('!I', socket.inet_aton("255.255.255.0"))[0]
         broadcast_bin = ip_bin | ~netmask_bin & 0xFFFFFFFF
         broadcast_ip = socket.inet_ntoa(struct.pack('!I', broadcast_bin))
         while self.running:
-            msg = f"SimPub:{json.dumps(local_info)}"
+            msg = f"SimPub:{_id}:{json.dumps(local_info)}"
             _socket.sendto(msg.encode(), (broadcast_ip, ServerPort.DISCOVERY))
             sleep(0.5)
         logger.info("Broadcasting has been stopped")

simpub/core/simpub_server.py

@@ -16,6 +16,8 @@ def __init__(self, host: str = "127.0.0.1"):
         self.host: str = host
         self.net_manager = init_net_manager(host)
 
+    def join(self):
+        self.net_manager.join()
 
 class MsgServer(ServerBase):
     pass

simpub/simdata.py

@@ -6,6 +6,8 @@
 import random
 import json
 
+# TODO: Using dict not dataclass to store for too many never-used attributes
+
 
 class VisualType(str, Enum):
     CUBE = "CUBE"

simpub/xr_device/xr_device.py

@@ -1,6 +1,10 @@
 import json
+import zmq
+from typing import Dict, Callable
+import time
 
-from ..core.subscriber import Subscriber, logger
+from ..core.subscriber import logger
+from ..core.net_manager import NetManager, Topic, ClientPort
 
 
 class InputData:
@@ -16,8 +20,55 @@ def __init__(
         self,
         device_name: str = "UnityEditor",
     ) -> None:
+        self.connected = False
+        self.manager: NetManager = NetManager.manager
         self.device = device_name
-        self.log_subscriber = Subscriber(f"{device_name}/Log", self.print_log)
+        self.client_info = None
+        # subscriber
+        self.sub_socket = self.manager.zmq_context.socket(zmq.SUB)
+        self.sub_topic_callback: Dict[Topic, Callable] = {}
+        self.register_topic_callback(f"{device_name}/Log", self.print_log)
+        # request client
+        self.req_socket = self.manager.zmq_context.socket(zmq.REQ)
+        self.manager.submit_task(self.wait_for_connection)
+
+    def wait_for_connection(self):
+        logger.info(f"Waiting for connection to {self.device}")
+        while not self.connected:
+            client_info = self.manager.clients_info.get(self.device)
+            if client_info is not None:
+                print(client_info)
+                self.connected = True
+                self.client_info = client_info
+                logger.info(f"Connected to {self.device}")
+                break
+            time.sleep(0.05)
+        self.req_socket.connect(
+            f"tcp://{self.client_info['ip']}:{ClientPort.SERVICE}")
+        self.sub_socket.connect(
+            f"tcp://{self.client_info['ip']}:{ClientPort.TOPIC}")
+        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "")
+        self.manager.submit_task(self.subscribe_loop)
+
+    def register_topic_callback(self, topic: str, callback: Callable):
+        self.sub_topic_callback[topic] = callback
+        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, topic)
+
+    def request(self, service: str, req: str) -> str:
+        if self.client_info is None:
+            logger.error(f"Device {self.device} is not connected")
+            return ""
+        if service not in self.client_info["services"]:
+            logger.error(f"Service {service} is not available")
+            return ""
+        self.req_socket.send_string(f"{service}:{req}")
+        return self.req_socket.recv_string()
+
+    def subscribe_loop(self):
+        while self.connected:
+            message = self.sub_socket.recv_string()
+            topic, msg = message.split(":", 1)
+            self.sub_topic_callback[topic](msg)
 
     def print_log(self, log: str):
         logger.info(f"{self.type} Log: {log}")