mcp_screen/test_stream.py

#!/usr/bin/env python3
import socket
import time
import argparse
import sys
from PIL import Image, ImageDraw, ImageFont

# Dimensions of the reflective LCD
WIDTH = 400
HEIGHT = 300

def map_to_rlcd_hw_buffer(img_1bit):
    """Converts a standard 1-bit monochrome PIL Image to the Waveshare RLCD hardware buffer layout."""
    px = img_1bit.load()
    hw_buffer = bytearray(15000)
    
    for y in range(HEIGHT):
        for x in range(WIDTH):
            if px[x, y]: # True if pixel is white (1)
                inv_y = HEIGHT - 1 - y
                byte_x  = x // 2
                block_y = inv_y // 4
                index = byte_x * 75 + block_y
                local_x = x % 2
                local_y = inv_y % 4
                bit = 7 - (local_y * 2 + local_x)
                hw_buffer[index] |= (1 << bit)
                
    return hw_buffer

def send_udp_frame_chunked(sock, esp32_ip, port, frame_idx, raw_bytes):
    """Sends a 15,000-byte frame split into 15 small UDP packets to bypass MTU and OS limits."""
    frame_id = frame_idx % 256
    for chunk_idx in range(15):
        packet = bytearray(1002)
        packet[0] = frame_id
        packet[1] = chunk_idx
        start = chunk_idx * 1000
        packet[2:1002] = raw_bytes[start : start + 1000]
        sock.sendto(packet, (esp32_ip, port))
        time.sleep(0.010) # Add small pacing delay to prevent network buffer flooding

def generate_animation_frame(frame_idx, mode_name):
    """Generates a test pattern animation frame using Pillow."""
    # Create 1-bit image (0 = black background)
    img = Image.new("1", (WIDTH, HEIGHT), 0)
    draw = ImageDraw.Draw(img)
    
    # Draw header text
    draw.text((10, 10), "ESP32-S3 VIDEO STREAM TEST", fill=1)
    draw.line((10, 22, 390, 22), fill=1)
    
    # Display active settings
    draw.text((15, 30), f"Protocol : {mode_name}", fill=1)
    draw.text((15, 45), f"Frame No : {frame_idx}", fill=1)
    
    # 1. Draw a bouncing rectangle
    rect_w, rect_h = 60, 40
    bounce_x = int((frame_idx * 5) % (WIDTH - rect_w - 20)) + 10
    bounce_y = int(120 + 20 * (frame_idx % 10 < 5 and (frame_idx % 5) or (5 - frame_idx % 5)))
    draw.rectangle([bounce_x, bounce_y, bounce_x + rect_w, bounce_y + rect_h], outline=1, width=2)
    draw.text((bounce_x + 10, bounce_y + 15), "TCP/UDP", fill=1)
    
    # 2. Draw a spinning needle / line
    import math
    angle = frame_idx * 0.1
    center_x, center_y = 300, 200
    radius = 40
    end_x = center_x + radius * math.cos(angle)
    end_y = center_y + radius * math.sin(angle)
    draw.ellipse([center_x - radius, center_y - radius, center_x + radius, center_y + radius], outline=1)
    draw.line((center_x, center_y, end_x, end_y), fill=1)
    
    # 3. Dynamic scrolling text at the bottom
    scrolling_text = "Waveshare RLCD 4.2inch -- Low Power -- High Framerate Streaming in MicroPython"
    scroll_pos = (frame_idx * 3) % 400
    draw.text((10 - scroll_pos, 275), scrolling_text, fill=1)
    draw.text((410 - scroll_pos, 275), scrolling_text, fill=1)
    draw.line((10, 270, 390, 270), fill=1)
    
    return img

def connect_socket(esp32_ip, port, protocol):
    """Establishes connection to the ESP32 stream server, retrying on failure for TCP."""
    while True:
        try:
            if protocol == "tcp":
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.connect((esp32_ip, port))
                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                print(f"Connected to stream server at {esp32_ip}:{port}")
                return sock
            else: # udp
                sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                if esp32_ip.endswith(".255") or esp32_ip == "255.255.255.255":
                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
                return sock
        except (socket.error, Exception) as e:
            if protocol == "tcp":
                print(f"Failed to connect to stream server: {e}. Retrying in 2 seconds...")
                time.sleep(2.0)
            else:
                # UDP doesn't establish connection, so any error here is fatal
                raise e

def stream_camera(esp32_ip, port, protocol):
    """Streams camera capture frames using OpenCV (requires opencv-python)."""
    try:
        import cv2
    except ImportError:
        print("Error: opencv-python is not installed. Run 'pip install opencv-python' or use default animation mode.")
        sys.exit(1)
        
    print(f"Opening camera stream...")
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Error: Could not open camera.")
        sys.exit(1)
        
    sock = connect_socket(esp32_ip, port, protocol)
        
    print(f"Streaming webcam to {esp32_ip}:{port} via {protocol.upper()}...")
    frame_count = 0
    start_time = time.time()
    
    try:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            # Flip image, convert to RGB
            frame = cv2.flip(frame, 1)
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_img = Image.fromarray(frame_rgb)
            
            # Resize and dither
            pil_img.thumbnail((WIDTH, HEIGHT))
            bg = Image.new("1", (WIDTH, HEIGHT), 0)
            offset = ((WIDTH - pil_img.width) // 2, (HEIGHT - pil_img.height) // 2)
            bg.paste(pil_img.convert("1", dither=Image.FLOYDSTEINBERG), offset)
            
            # Map pixels to hardware buffer
            raw_bytes = map_to_rlcd_hw_buffer(bg)
            
            # Transmit with reconnect logic
            try:
                if protocol == "tcp":
                    sock.sendall(raw_bytes)
                else:
                    send_udp_frame_chunked(sock, esp32_ip, port, frame_count, raw_bytes)
            except (socket.error, Exception) as se:
                print(f"Transmission failed: {se}. Reconnecting...")
                try:
                    sock.close()
                except:
                    pass
                sock = connect_socket(esp32_ip, port, protocol)
                
            frame_count += 1
            elapsed = time.time() - start_time
            if elapsed >= 2.0:
                print(f"Streaming performance: {frame_count / elapsed:.1f} FPS")
                frame_count = 0
                start_time = time.time()
                
            # Target ~30 FPS
            time.sleep(0.033)
            
    finally:
        cap.release()
        if sock:
            sock.close()

def stream_animation(esp32_ip, port, protocol):
    """Streams a generated Pillow vector animation."""
    sock = connect_socket(esp32_ip, port, protocol)
            
    print(f"Streaming vector animation to {esp32_ip}:{port} via {protocol.upper()}...")
    frame_idx = 0
    start_time = time.time()
    
    try:
        while True:
            # Generate frame
            pil_img = generate_animation_frame(frame_idx, f"{protocol.upper()} Stream")
            
            # Map pixels to hardware buffer
            raw_bytes = map_to_rlcd_hw_buffer(pil_img)
            
            # Transmit with reconnect logic
            try:
                if protocol == "tcp":
                    sock.sendall(raw_bytes)
                else:
                    send_udp_frame_chunked(sock, esp32_ip, port, frame_idx, raw_bytes)
            except (socket.error, Exception) as se:
                print(f"Transmission failed: {se}. Reconnecting...")
                try:
                    sock.close()
                except:
                    pass
                sock = connect_socket(esp32_ip, port, protocol)
                
            frame_idx += 1
            
            # FPS tracking printout
            if frame_idx % 60 == 0:
                elapsed = time.time() - start_time
                print(f"Streaming performance: {60 / elapsed:.1f} FPS")
                start_time = time.time()
                
            # Restrict rate to ~25 FPS to match update scan limits nicely
            time.sleep(0.04)
            
    except KeyboardInterrupt:
        print("\nStreaming stopped.")
    finally:
        if sock:
            sock.close()

def main():
    parser = argparse.ArgumentParser(description="Host Client for ESP32-S3 TCP/UDP Video Streaming")
    parser.add_argument("--ip", default="192.168.68.123", help="Target ESP32-S3 IP Address (default: 192.168.68.123)")
    parser.add_argument("--protocol", choices=["tcp", "udp"], default="tcp", help="Streaming protocol (default: tcp)")
    parser.add_argument("--source", choices=["camera", "animation"], default="animation", help="Streaming source (default: animation)")
    args = parser.parse_args()
    
    # Assign ports based on protocol selection
    port = 8081 if args.protocol == "tcp" else 8082
    
    if args.source == "camera":
        stream_camera(args.ip, port, args.protocol)
    else:
        stream_animation(args.ip, port, args.protocol)

if __name__ == "__main__":
    main()