282 lines
8.9 KiB
Python
282 lines
8.9 KiB
Python
"""UDP server for receiving camera frames from ESP devices.
|
|
|
|
Protocol from ESP32:
|
|
- TOKEN + "START" -> Start of new frame
|
|
- TOKEN + chunk -> JPEG data chunk
|
|
- TOKEN + "END" -> End of frame, decode and process
|
|
"""
|
|
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
import cv2
|
|
import numpy as np
|
|
from typing import Optional, Callable, Dict
|
|
|
|
from .config import (
|
|
UDP_HOST, UDP_PORT, UDP_BUFFER_SIZE,
|
|
SECRET_TOKEN, IMAGE_DIR,
|
|
VIDEO_ENABLED, VIDEO_PATH, VIDEO_FPS, VIDEO_CODEC
|
|
)
|
|
|
|
|
|
class FrameAssembler:
|
|
"""Assembles JPEG frames from multiple UDP packets."""
|
|
|
|
def __init__(self, timeout: float = 5.0):
|
|
self.timeout = timeout
|
|
self.buffer = bytearray()
|
|
self.start_time: Optional[float] = None
|
|
self.receiving = False
|
|
|
|
def start_frame(self):
|
|
"""Start receiving a new frame."""
|
|
self.buffer = bytearray()
|
|
self.start_time = time.time()
|
|
self.receiving = True
|
|
|
|
def add_chunk(self, data: bytes) -> bool:
|
|
"""Add a chunk to the frame buffer. Returns False if timed out."""
|
|
if not self.receiving:
|
|
return False
|
|
if self.start_time and (time.time() - self.start_time) > self.timeout:
|
|
self.reset()
|
|
return False
|
|
self.buffer.extend(data)
|
|
return True
|
|
|
|
def finish_frame(self) -> Optional[bytes]:
|
|
"""Finish frame assembly and return complete data."""
|
|
if not self.receiving or len(self.buffer) == 0:
|
|
return None
|
|
data = bytes(self.buffer)
|
|
self.reset()
|
|
return data
|
|
|
|
def reset(self):
|
|
"""Reset the assembler state."""
|
|
self.buffer = bytearray()
|
|
self.start_time = None
|
|
self.receiving = False
|
|
|
|
|
|
class UDPReceiver:
|
|
"""Receives JPEG frames via UDP from ESP camera devices."""
|
|
|
|
def __init__(self,
|
|
host: str = UDP_HOST,
|
|
port: int = UDP_PORT,
|
|
image_dir: str = IMAGE_DIR,
|
|
on_frame: Optional[Callable] = None):
|
|
self.host = host
|
|
self.port = port
|
|
self.image_dir = image_dir
|
|
self.on_frame = on_frame # Callback when frame received
|
|
|
|
self._sock: Optional[socket.socket] = None
|
|
self._thread: Optional[threading.Thread] = None
|
|
self._stop_event = threading.Event()
|
|
|
|
# Frame assemblers per source address
|
|
self._assemblers: Dict[str, FrameAssembler] = {}
|
|
|
|
# Video recording
|
|
self._video_writer: Optional[cv2.VideoWriter] = None
|
|
self._video_size: Optional[tuple] = None
|
|
|
|
# Statistics
|
|
self.frames_received = 0
|
|
self.invalid_tokens = 0
|
|
self.decode_errors = 0
|
|
self.packets_received = 0
|
|
|
|
# Active cameras tracking
|
|
self._active_cameras: dict = {} # {camera_id: last_frame_time}
|
|
|
|
os.makedirs(self.image_dir, exist_ok=True)
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._thread is not None and self._thread.is_alive()
|
|
|
|
@property
|
|
def active_cameras(self) -> list:
|
|
"""Returns list of active camera identifiers."""
|
|
return list(self._active_cameras.keys())
|
|
|
|
def start(self) -> bool:
|
|
"""Start the UDP receiver thread."""
|
|
if self.is_running:
|
|
return False
|
|
|
|
self._stop_event.clear()
|
|
self._thread = threading.Thread(target=self._receive_loop, daemon=True)
|
|
self._thread.start()
|
|
return True
|
|
|
|
def stop(self):
|
|
"""Stop the UDP receiver and cleanup."""
|
|
self._stop_event.set()
|
|
|
|
if self._sock:
|
|
try:
|
|
self._sock.close()
|
|
except Exception:
|
|
pass
|
|
self._sock = None
|
|
|
|
if self._video_writer is not None:
|
|
self._video_writer.release()
|
|
self._video_writer = None
|
|
|
|
# Clean up frame files
|
|
self._cleanup_frames()
|
|
|
|
self._active_cameras.clear()
|
|
self._assemblers.clear()
|
|
self.frames_received = 0
|
|
self.packets_received = 0
|
|
|
|
def _cleanup_frames(self):
|
|
"""Remove all .jpg files from image directory."""
|
|
try:
|
|
for f in os.listdir(self.image_dir):
|
|
if f.endswith(".jpg"):
|
|
os.remove(os.path.join(self.image_dir, f))
|
|
except Exception:
|
|
pass
|
|
|
|
def _get_assembler(self, addr: tuple) -> FrameAssembler:
|
|
"""Get or create a frame assembler for the given address."""
|
|
key = f"{addr[0]}:{addr[1]}"
|
|
if key not in self._assemblers:
|
|
self._assemblers[key] = FrameAssembler()
|
|
return self._assemblers[key]
|
|
|
|
def _receive_loop(self):
|
|
"""Main UDP receive loop with START/END protocol handling."""
|
|
self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
self._sock.bind((self.host, self.port))
|
|
self._sock.settimeout(1.0)
|
|
|
|
print(f"[UDP] Receiver started on {self.host}:{self.port}")
|
|
|
|
while not self._stop_event.is_set():
|
|
try:
|
|
data, addr = self._sock.recvfrom(UDP_BUFFER_SIZE)
|
|
except socket.timeout:
|
|
continue
|
|
except OSError:
|
|
break
|
|
|
|
self.packets_received += 1
|
|
|
|
# Validate token
|
|
if not data.startswith(SECRET_TOKEN):
|
|
self.invalid_tokens += 1
|
|
continue
|
|
|
|
# Remove token prefix
|
|
payload = data[len(SECRET_TOKEN):]
|
|
assembler = self._get_assembler(addr)
|
|
camera_id = f"{addr[0]}_{addr[1]}"
|
|
|
|
# Handle protocol markers
|
|
if payload == b"START":
|
|
assembler.start_frame()
|
|
continue
|
|
elif payload == b"END":
|
|
frame_data = assembler.finish_frame()
|
|
if frame_data:
|
|
self._process_complete_frame(camera_id, frame_data, addr)
|
|
continue
|
|
else:
|
|
# Regular data chunk
|
|
if not assembler.receiving:
|
|
# No START received, try as single-packet frame (legacy)
|
|
frame = self._decode_frame(payload)
|
|
if frame is not None:
|
|
self._process_frame(camera_id, frame, addr)
|
|
else:
|
|
self.decode_errors += 1
|
|
else:
|
|
assembler.add_chunk(payload)
|
|
|
|
# Cleanup
|
|
if self._sock:
|
|
self._sock.close()
|
|
self._sock = None
|
|
|
|
if self._video_writer:
|
|
self._video_writer.release()
|
|
self._video_writer = None
|
|
|
|
print("[UDP] Receiver stopped")
|
|
|
|
def _process_complete_frame(self, camera_id: str, frame_data: bytes, addr: tuple):
|
|
"""Process a fully assembled frame."""
|
|
frame = self._decode_frame(frame_data)
|
|
if frame is None:
|
|
self.decode_errors += 1
|
|
return
|
|
self._process_frame(camera_id, frame, addr)
|
|
|
|
def _process_frame(self, camera_id: str, frame: np.ndarray, addr: tuple):
|
|
"""Process a decoded frame."""
|
|
self.frames_received += 1
|
|
self._active_cameras[camera_id] = time.time()
|
|
|
|
# Save frame
|
|
self._save_frame(camera_id, frame)
|
|
|
|
# Record video if enabled
|
|
if VIDEO_ENABLED:
|
|
self._record_frame(frame)
|
|
|
|
# Callback
|
|
if self.on_frame:
|
|
self.on_frame(camera_id, frame, addr)
|
|
|
|
def _decode_frame(self, data: bytes) -> Optional[np.ndarray]:
|
|
"""Decode JPEG data to OpenCV frame."""
|
|
try:
|
|
npdata = np.frombuffer(data, np.uint8)
|
|
frame = cv2.imdecode(npdata, cv2.IMREAD_COLOR)
|
|
return frame
|
|
except Exception:
|
|
return None
|
|
|
|
def _save_frame(self, camera_id: str, frame: np.ndarray):
|
|
"""Save frame as JPEG file."""
|
|
try:
|
|
filepath = os.path.join(self.image_dir, f"{camera_id}.jpg")
|
|
cv2.imwrite(filepath, frame)
|
|
except Exception:
|
|
pass
|
|
|
|
def _record_frame(self, frame: np.ndarray):
|
|
"""Record frame to video file."""
|
|
if self._video_writer is None:
|
|
self._video_size = (frame.shape[1], frame.shape[0])
|
|
fourcc = cv2.VideoWriter_fourcc(*VIDEO_CODEC)
|
|
video_path = os.path.join(os.path.dirname(self.image_dir), VIDEO_PATH.split('/')[-1])
|
|
self._video_writer = cv2.VideoWriter(
|
|
video_path, fourcc, VIDEO_FPS, self._video_size
|
|
)
|
|
|
|
if self._video_writer and self._video_writer.isOpened():
|
|
self._video_writer.write(frame)
|
|
|
|
def get_stats(self) -> dict:
|
|
"""Return receiver statistics."""
|
|
return {
|
|
"running": self.is_running,
|
|
"packets_received": self.packets_received,
|
|
"frames_received": self.frames_received,
|
|
"invalid_tokens": self.invalid_tokens,
|
|
"decode_errors": self.decode_errors,
|
|
"active_cameras": len(self._active_cameras)
|
|
}
|