Merge pull request #52 from AccelerationConsortium/cobotmqtt

added server and client files for cobot mqtt control

src/ac_training_lab/cobot280pi/.gitignore

@@ -0,0 +1,3 @@
+my_secrets.py
+*.jpg
+.stfolder/

src/ac_training_lab/cobot280pi/client.py

@@ -0,0 +1,100 @@
+import base64
+import io
+import json
+from queue import Queue
+
+import paho.mqtt.client as paho
+from PIL import Image
+
+
+class CobotController:
+
+    def __init__(
+        self,
+        hive_mq_username: str,
+        hive_mq_password: str,
+        hive_mq_cloud: str,
+        port: int,
+        cobot_id: str,
+    ):
+        self.publish_endpoint = cobot_id
+        self.response_endpoint = cobot_id + "/response"
+        self.client = paho.Client(client_id="", userdata=None, protocol=paho.MQTTv5)
+        self.client.tls_set()
+        self.client.username_pw_set(hive_mq_username, hive_mq_password)
+        self.client.connect(hive_mq_cloud, port)
+
+        response_queue = Queue()
+
+        def on_message(client, userdata, msg):
+            payload_dict = json.loads(msg.payload)
+            response_queue.put(payload_dict)
+
+        self.response_queue = response_queue
+
+        def on_connect(client, userdata, flags, rc, properties=None):
+            print("Connection recieved")
+
+        self.client.on_connect = on_connect
+        self.client.on_message = on_message
+        self.client.subscribe(self.response_endpoint, qos=2)
+        self.client.loop_start()
+
+    def handle_publish_and_response(self, payload):
+        self.client.publish(self.publish_endpoint, payload=payload, qos=2)
+        return self.response_queue.get(block=True)
+
+    def send_angles(self, angle_list: list[float] = [0.0] * 6, speed: int = 50):
+        payload = json.dumps(
+            {
+                "command": "control/angles",
+                "args": {"angles": angle_list, "speed": speed},
+            }
+        )
+        return self.handle_publish_and_response(payload)
+
+    def send_coords(self, coord_list: list[float] = [0.0] * 6, speed: int = 50):
+        payload = json.dumps(
+            {
+                "command": "control/coords",
+                "args": {"coords": coord_list, "speed": speed},
+            }
+        )
+        return self.handle_publish_and_response(payload)
+
+    def send_gripper_value(self, value: int = 100, speed: int = 50):
+        payload = json.dumps(
+            {
+                "command": "control/gripper",
+                "args": {"gripper_value": value, "speed": speed},
+            }
+        )
+        return self.handle_publish_and_response(payload)
+
+    def get_angles(self):
+        payload = json.dumps({"command": "query/angles", "args": {}})
+        return self.handle_publish_and_response(payload)
+
+    def get_coords(self):
+        payload = json.dumps({"command": "query/coords", "args": {}})
+        return self.handle_publish_and_response(payload)
+
+    def get_gripper_value(self):
+        payload = json.dumps({"command": "query/gripper", "args": {}})
+        return self.handle_publish_and_response(payload)
+
+    def get_camera(self, quality=100, save_path=None):
+        payload = json.dumps({"command": "query/camera", "args": {"quality": quality}})
+        response = self.handle_publish_and_response(payload)
+        if not response["success"]:
+            return response
+
+        b64_bytes = base64.b64decode(response["image"])
+        img_bytes = io.BytesIO(b64_bytes)
+        img = Image.open(img_bytes)
+
+        response["image"] = img
+        if save_path is not None:
+            img.save(save_path)
+
+        return response

src/ac_training_lab/cobot280pi/demo.py

@@ -0,0 +1,11 @@
+from pymycobot.mycobot import MyCobot
+
+cobot = MyCobot("/dev/ttyAMA0", 1000000)
+
+
+def rise():
+    cobot.send_angles([0, 0, 0, 0, 0, 0], 100)
+
+
+def kill():
+    cobot.release_all_servos()

src/ac_training_lab/cobot280pi/dummy_cobot.py

@@ -0,0 +1,34 @@
+from PIL import Image
+from utils import setup_logger
+
+
+# A dummy class for easier testing without physically having the cobot
+class DummyCobot:
+
+    def __init__(self):
+        self.logger = setup_logger()
+
+    def set_gripper_value(self, **kwargs):
+        self.logger.info(f"tried to set gripper value with args {kwargs}")
+
+    def send_angles(self, **kwargs):
+        self.logger.info(f"tried to send angles with args {kwargs}")
+
+    def send_coords(self, **kwargs):
+        self.logger.info(f"tried to send coords with args {kwargs}")
+
+    def get_angles(self, **kwargs):
+        self.logger.info(f"tried to get angles with args {kwargs}")
+        return [0, 0, 0, 0, 0, 0]
+
+    def get_coords(self, **kwargs):
+        self.logger.info(f"tried to get coords with args {kwargs}")
+        return [0, 0, 0, 0, 0, 0]
+
+    def get_gripper_value(self, **kwargs):
+        self.logger.info(f"tried to get gripper value with args {kwargs}")
+        return 0
+
+    def get_camera(self, **kwargs):
+        self.logger.info(f"tried to get camera with args {kwargs}")
+        return Image.new("RGB", (1920, 1080), color="black")

src/ac_training_lab/cobot280pi/requirements.txt

@@ -0,0 +1,4 @@
+paho-mqtt==2.1.0
+pillow==10.4.0
+setuptools==75.1.0
+wheel==0.44.0

src/ac_training_lab/cobot280pi/server.py

@@ -0,0 +1,194 @@
+import argparse
+import base64
+import io
+import json
+import sys
+import time
+from queue import Queue
+
+import cv2
+import paho.mqtt.client as paho
+from my_secrets import (
+    DEVICE_ENDPOINT,
+    DEVICE_PORT,
+    HIVEMQ_HOST,
+    HIVEMQ_PASSWORD,
+    HIVEMQ_USERNAME,
+)
+from PIL import Image
+from pymycobot.mycobot import MyCobot
+from utils import setup_logger
+
+# cli args
+parser = argparse.ArgumentParser()
+parser.add_argument("--debug", "-d", action="store_true", help="runs in debug mode")
+cliargs = parser.parse_args()
+
+
+# Cobot action functions
+def handle_control_gripper(args, cobot):
+    logger.info(f"running command control/gripper with {args}")
+    try:
+        cobot.set_gripper_value(**args)
+        return {"success": True}
+    except Exception as e:
+        logger.critical(f"control gripper error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+def handle_control_angles(args, cobot):
+    logger.info(f"running command control/angle with {args}")
+    try:
+        cobot.send_angles(**args)
+        return {"success": True}
+    except Exception as e:
+        logger.critical(f"control angle error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+def handle_control_coords(args, cobot):
+    logger.info(f"running command control/coord with {args}")
+    try:
+        cobot.send_coords(**args)
+        return {"success": True}
+    except Exception as e:
+        logger.critical(f"control coords error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+def handle_query_angles(args, cobot):
+    logger.info(f"running command query/angle with {args}")
+    try:
+        angles = cobot.get_angles()
+        if angles is None or len(angles) < 6:
+            raise Exception("could not read angle")
+        return {"success": True, "angles": angles}
+    except Exception as e:
+        logger.critical(f"query angle error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+def handle_query_coords(args, cobot):
+    logger.info(f"running command query/coord with {args}")
+    try:
+        coords = cobot.get_coords()
+        if coords is None or len(coords) < 6:
+            raise Exception("could not read coord")
+        return {"success": True, "coords": coords}
+    except Exception as e:
+        logger.critical(f"query coord error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+def handle_query_gripper(args, cobot):
+    logger.info(f"running command query/coord with {args}")
+    try:
+        gripper_pos = cobot.get_gripper_value()
+        return {"success": True, "position": gripper_pos}
+    except Exception as e:
+        logger.critical(f"query gripper error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+def handle_query_camera(args):
+    logger.info(f"running command query/camera with {args}")
+    try:
+        if not cliargs.debug:
+            webcam = cv2.VideoCapture(0)
+            _, frame = webcam.read()
+            webcam.release()
+            img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        else:
+            img = cobot.get_camera(**args)
+        compressed_bytes = io.BytesIO()
+        img.save(compressed_bytes, format="JPEG", quality=args["quality"])
+        compressed_bytes.seek(0)
+        byte_str = base64.b64encode(compressed_bytes.read()).decode("utf-8")
+        return {"success": True, "image": byte_str}
+    except Exception as e:
+        logger.critical(f"query camera error: {str(e)}")
+        return {"success": False, "error_msg": str(e)}
+
+
+# MQTT Functions
+def on_connect(client, userdata, flags, rc, properties=None):
+    logger.info("Connection received with code %s." % rc)
+
+
+def on_publish(client, userdata, mid, properties=None):
+    logger.info("Successful publish.")
+
+
+def handle_message(msg, cobot):
+    # Parse payload to json dict
+    try:
+        payload_dict = json.loads(msg.payload)
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+    if "command" not in payload_dict:
+        return {"success": False, "error": "'command' key should be in payload"}
+
+    # Match to a command function
+    cmd = payload_dict["command"]
+    if cmd == "control/angles":
+        return handle_control_angles(payload_dict["args"], cobot)
+    elif cmd == "control/coords":
+        return handle_control_coords(payload_dict["args"], cobot)
+    elif cmd == "control/gripper":
+        return handle_control_gripper(payload_dict["args"], cobot)
+    elif cmd == "query/angles":
+        return handle_query_angles(payload_dict["args"], cobot)
+    elif cmd == "query/coords":
+        return handle_query_coords(payload_dict["args"], cobot)
+    elif cmd == "query/gripper":
+        return handle_query_gripper(payload_dict["args"], cobot)
+    elif cmd == "query/camera":
+        return handle_query_camera(payload_dict["args"])
+    else:
+        return {"success": False, "error": "invalid command"}
+
+
+if __name__ == "__main__":
+    task_queue = Queue()
+    logger = setup_logger()
+
+    if not cliargs.debug:
+        try:
+            cobot = MyCobot("/dev/ttyAMA0", 1000000)
+            logger.info("Cobot object initialized...")
+        except Exception as e:
+            logger.critical(f"could not initialize cobot with error {str(e)}")
+            sys.exit(1)
+    else:
+        from dummy_cobot import DummyCobot
+
+        cobot = DummyCobot()
+
+    def on_message(client, userdata, msg):
+        logger.info(
+            f"Recieved message with: \n\ttopic: {msg.topic}\
+            \n\tqos: {msg.qos}\n\tpayload: {msg.payload}"
+        )
+        task_queue.put(msg)
+
+    client = paho.Client(client_id="", userdata=None, protocol=paho.MQTTv5)
+    client.on_connect = on_connect
+    client.on_message = on_message
+    client.on_publish = on_publish
+
+    client.tls_set(tls_version=paho.ssl.PROTOCOL_TLS)
+    client.username_pw_set(HIVEMQ_USERNAME, HIVEMQ_PASSWORD)
+    client.connect(HIVEMQ_HOST, DEVICE_PORT)
+    client.subscribe(DEVICE_ENDPOINT, qos=2)
+    client.loop_start()
+    logger.info("Ready for tasks...")
+
+    while True:
+        msg = task_queue.get()  # blocks if queue is empty
+        response_dict = handle_message(msg, cobot)
+        pub_handle = client.publish(
+            DEVICE_ENDPOINT + "/response", qos=2, payload=json.dumps(response_dict)
+        )
+        pub_handle.wait_for_publish()
+        time.sleep(3)

src/ac_training_lab/cobot280pi/utils.py

@@ -0,0 +1,20 @@
+import logging
+import sys
+
+
+def setup_logger(logfile_name: str = "mqttcobot.log"):
+    logger = logging.getLogger("logger")
+    logger.setLevel(logging.INFO)
+
+    console_handler = logging.StreamHandler(sys.stdout)
+    file_handler = logging.FileHandler(logfile_name)
+
+    formatter = logging.Formatter("[%(levelname)s - %(asctime)s]: %(message)s")
+    console_handler.setFormatter(formatter)
+    file_handler.setFormatter(formatter)
+
+    if not logger.hasHandlers():
+        logger.addHandler(console_handler)
+        logger.addHandler(file_handler)
+    logger.propagate = False
+    return logger