add nicer console logging to server

README.md

@@ -29,12 +29,30 @@ source .venv/bin/activate
 
 ### Client
 
+To run the client with a localhost target:
+
 ```shell
 python -m client
 ```
 
+To target an external server, provide its IP address using the `-s` flag:
+
+```shell
+python -m client -s [server IP address]
+```
+
 ### Server
 
 ```shell
 python -m server
-```
+```
+
+### Common Flags
+
+Make sure that these match between your client and server!!
+
+| Short | Long | Description | Default |
+|-------|---|---|---|
+| `-p`  | `--port` | The port for the client / server to communicate on. | `5005` |
+| `-W`  | `--width` | Image width in pixels. | `640` |
+| `-H` | `--height` | Image height in pixels. | `480` |

client.py

@@ -1,28 +1,15 @@
+import argparse
+
 import cv2
 import socket
 import numpy as np
 import uuid
 
-uuid = uuid.uuid4()
-
 from common import StdPacket, InterlacedPacket, DoublyInterlacedPacket
 
-UDP_IP = "127.0.0.1"
-UDP_PORT = 5005
-
-# Create a VideoCapture object
-cap = cv2.VideoCapture(0)  # 0 represents the default camera
-
-# Check if camera opened successfully
-if not cap.isOpened():
-    print("Error opening video stream or file")
-
 def send_packet(sock, packet):
     sock.sendto(packet, (UDP_IP, UDP_PORT))
 
-sock = socket.socket(socket.AF_INET, # Internet
-                     socket.SOCK_DGRAM) # UDP
-
 def breakdown_image_norm(frame):
     (cols, rows, colors) = frame.shape
     # break the array down into 16x16 chunks, then transmit them as UDP packets
@@ -61,16 +48,48 @@ def breakdown_image_dint(frame):
                 pkt = DoublyInterlacedPacket(uuid, j, i, j_even, i_even, frame[i + i_even:i + 32:2, j + j_even:j + 32:2])
                 send_packet(sock, pkt.to_bytestr())
 
-while True:
+if __name__ == '__main__':
-    # Capture frame-by-frame
+    # argument parser
-    ret, frame = cap.read()
+    parser = argparse.ArgumentParser(description="Proof-of-concept client for sauron-cv")
+    parser.add_argument("-p", "--port", type=int, default=5005)
+    parser.add_argument("-s", "--server", type=str, default="127.0.0.1")
+    parser.add_argument("-W", "--width", type=int, default=640)
+    parser.add_argument("-H", "--height", type=int, default=480)
+    parser.add_argument("-d", "--device", type=int, default=0)
+    args = parser.parse_args()
 
-    # If frame is read correctly, ret is True
+    # give this client its very own UUID
-    if not ret:
+    uuid = uuid.uuid4()
-        print("Can't receive frame (stream end?). Exiting ...")
-        break
 
-    breakdown_image_dint(frame)
+    # give target IP address
+    UDP_IP = args.server
+    UDP_PORT = args.port
+    WIDTH = args.width
+    HEIGHT = args.height
+    DEVICE = args.device
 
-# Release the capture and close all windows
+    # Create a VideoCapture object
-cap.release()
+    cap = cv2.VideoCapture(DEVICE)  # 0 represents the default camera
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)
+
+    # Check if camera opened successfully
+    if not cap.isOpened():
+        print("Error opening video stream or file")
+
+    # create the socket
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+
+    while True:
+        # Capture frame-by-frame
+        ret, frame = cap.read()
+
+        # If frame is read correctly, ret is True
+        if not ret:
+            print("Can't receive frame (stream end?). Exiting ...")
+            break
+
+        breakdown_image_dint(frame)
+
+    # Release the capture and close all windows
+    cap.release()

common.py

@@ -3,7 +3,18 @@ from abc import ABC, abstractmethod
 import numpy as np
 from uuid import UUID
 
+# The basic structure of the packet is as follows:
+#  FFFF FFFF FFFF FFFF FFFF FFFF FFFF FFFF FFFF FFFF FFFF FFFF FFFF ... FFFF
+# | uuid              |    |    |                                           |
+# |                   |    |    | image data: contains the 16x16 image slice bitpacked from a NumPy array
+# |                   |    | y: the y position of this packet in the original image
+# |                   | x: the x position of this packet in the original image
+# | uuid: matches the packet to the requesting client
+# Other packet types may change this structure. Need to standardize this somehow.
+
 class Packet(ABC):
+    """Generic structure for a video streaming packet. Contains a slice of the full image, which will be reconstructed
+    by the server."""
     size: int
     def __init__(self, uuid: UUID, x: int, y: int, array: np.ndarray):
         self.uuid = uuid
@@ -13,13 +24,16 @@ class Packet(ABC):
 
     @abstractmethod
     def to_bytestr(self) -> bytes:
+        """Convert a packet object into a bytestring."""
         pass
 
     @abstractmethod
     def apply(self, image: np.ndarray) -> np.ndarray:
+        """Apply this packet to an image."""
         pass
 
 class StdPacket(Packet):
+    """A standard packet with no interlacing. Sends a 16x16 chunk of an image."""
     size = 16 + 4 + 4 + 768
     def __init__(self, uuid: UUID, x: int, y: int, array: np.ndarray):
         super().__init__(uuid, x, y, array)
@@ -40,6 +54,7 @@ class StdPacket(Packet):
         return image
 
 def from_bytes_std(b: bytes) -> StdPacket:
+    """Convert a byte string obtained via UDP into a packet object."""
     uuid = UUID(bytes = b[0:16])
     x = int.from_bytes(b[16:20], signed = False)
     y = int.from_bytes(b[20:24], signed = False)
@@ -49,6 +64,7 @@ def from_bytes_std(b: bytes) -> StdPacket:
 
 
 class InterlacedPacket(Packet):
+    """A packet with horizontal interlacing. Sends half of a 16x32 chunk of an image, every other row"""
     size = 16 + 4 + 4 + 4 + 768
     def __init__(self, uuid: UUID, x: int, y: int, even: bool, array: np.ndarray):
         super().__init__(uuid, x, y, array)
@@ -72,6 +88,7 @@ class InterlacedPacket(Packet):
 
 
 def from_bytes_int(b: bytes) -> InterlacedPacket:
+    """Convert a byte string obtained via UDP into a packet object."""
     uuid = UUID(bytes=b[0:16])
     x = int.from_bytes(b[16:20], signed=False)
     y = int.from_bytes(b[20:24], signed=False)
@@ -81,6 +98,8 @@ def from_bytes_int(b: bytes) -> InterlacedPacket:
     return InterlacedPacket(uuid, x, y, even, array)
 
 class DoublyInterlacedPacket(Packet):
+    """A packet with horizontal interlacing. Sends one quarter of a 32x32 chunk of an image. This will alternate rows
+    and columns based on the value of even_x and even_y."""
     size = 16 + 4 + 4 + 4 + 768
     def __init__(self, uuid: UUID, x: int, y: int, even_x: bool, even_y: bool, array: np.ndarray):
         super().__init__(uuid, x, y, array)
@@ -106,6 +125,7 @@ class DoublyInterlacedPacket(Packet):
 
 
 def from_bytes_dint(b: bytes) -> DoublyInterlacedPacket:
+    """Convert a byte string obtained via UDP into a packet object."""
     uuid = UUID(bytes=b[0:16])
     x = int.from_bytes(b[16:20], signed=False)
     y = int.from_bytes(b[20:24], signed=False)

requirements.txt

@@ -1,2 +1,3 @@
 opencv-python
-numpy
+numpy
+rich

server.py

@@ -5,87 +5,121 @@ import socket
 import numpy as np
 from datetime import datetime
 import asyncio
+import argparse
+from rich.console import Console
 
-from common import DoublyInterlacedPacket, from_bytes_dint
+from common import Packet, DoublyInterlacedPacket, from_bytes_dint
-
-# bind any IP address
-UDP_IP = ""
-UDP_PORT = 5005
-
-HEIGHT = 480
-WIDTH = 640
 
 class Client:
+    """Class for tracking client state, including current frame data and time since last update."""
     def __init__(self):
         self.last_updated = datetime.now()
         self.frame = np.ndarray((HEIGHT, WIDTH, 3), dtype=np.uint8)
 
-    def update(self, pkt: DoublyInterlacedPacket):
+    def update(self, pkt: Packet):
+        """Apply a packet to the client frame. Update last client update to current time."""
         self.frame = pkt.apply(self.frame)
 
         self.last_updated = datetime.now()
 
     def latency(self) -> float:
+        """Return the time since the last client update."""
         return (datetime.now() - self.last_updated).total_seconds()
 
     def read(self) -> np.ndarray:
+        """Return the current frame."""
         return self.frame
 
 
+# Dictionary of client states stored by UUID
 frames: Dict[str, Client] = {}
 
 async def read_packet():
+    """Asynchronous coroutine to read UDP packets from the client(s)."""
     while True:
-        for i in range(0, 1200):
+        # we repeat this a ton of times all at once to hopefully capture all of the image data
-        # break the array down into 16-bit chunks, then transmit them as UDP packets
+        for i in range(0, 1600):
             try:
-                data, addr = sock.recvfrom(DoublyInterlacedPacket.size)  # buffer size is 768 bytes
+                data, addr = sock.recvfrom(DoublyInterlacedPacket.size)  # packet buffer size based on the packet size
                 # print("received packet from", addr)
                 if data:
                     pkt = from_bytes_dint(data)
 
                     uuid = str(pkt.uuid)
 
+                    # if this is a new client, give it a new image
                     if uuid not in frames.keys():
-                        print("New client acquired, naming %s", uuid)
+                        console.log(f"New client acquired, naming [bold cyan]{uuid}[bold cyan]")
                         frames[uuid] = Client()
+                        stat.update(f"[bold yellow]{len(frames.keys())}[/bold yellow] clients connected.")
 
                     frames[uuid].update(pkt)
 
             except BlockingIOError:
                 pass
 
+        # this is necessary to allow asyncio to swap between reading packets and rendering frames
         await asyncio.sleep(0.001)
 
 async def show_frames():
+    """Asynchronous coroutine to display frames in OpenCV debug windows."""
     while True:
-        # Display the resulting frame
+        # drop clients that have not sent packets for > 5 seconds
         for id in list(frames.keys()):
             if frames[id].latency() >= 5:
-                print("Client likely lost connection, dropping %s", id)
+                console.log(f"Client likely lost connection, dropping [bold red]{id}[/bold red]")
                 cv2.destroyWindow(id)
                 frames.pop(id)
+                stat.update(f"[bold yellow]{len(frames.keys())}[/bold yellow] clients connected.")
             else:
+                # show the latest available frame
                 cv2.imshow(id, frames[id].read())
 
         cv2.waitKey(1)
+        # this is necessary to allow asyncio to swap between reading packets and rendering frames
         await asyncio.sleep(0.01)
 
 if __name__ == "__main__":
+    # argument parser
+    parser = argparse.ArgumentParser(description="Proof-of-concept server for sauron-cv")
+    parser.add_argument("-p", "--port", type=int, default=5005)
+    parser.add_argument("-l", "--listen", type=str, default="")
+    parser.add_argument("-W", "--width", type=int, default=640)
+    parser.add_argument("-H", "--height", type=int, default=480)
+    args = parser.parse_args()
+
+    # console
+    console = Console()
+
+    # assign constants based on argument parser
+    UDP_IP = args.listen
+    UDP_PORT = args.port
+
+    HEIGHT = args.height
+    WIDTH = args.width
+
+    # create the UDP socket
     sock = socket.socket(socket.AF_INET,  # Internet
                          socket.SOCK_DGRAM)  # UDP
     sock.setblocking(False)
     sock.bind((UDP_IP, UDP_PORT))
 
-    loop = asyncio.new_event_loop()
+    console.log("Ready to accept connections.", style="bold green")
-    asyncio.set_event_loop(loop)
-    loop.create_task(read_packet())
-    loop.create_task(show_frames())
-    try:
-        loop.run_forever()
-    finally:
-        loop.run_until_complete(loop.shutdown_asyncgens())
-        loop.close()
 
-        # Release the capture and close all windows
+    with console.status("[bold yellow]0[/bold yellow] clients connected.", spinner="pong") as stat:
-        cv2.destroyAllWindows()
+
+        # create the async event loop
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        # create async tasks for reading network packets, displaying windows
+        loop.create_task(read_packet())
+        loop.create_task(show_frames())
+        try:
+            loop.run_forever()
+        finally:
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            loop.close()
+
+            # Release the capture and close all windows
+            cv2.destroyAllWindows()