469 lines
15 KiB
Python
469 lines
15 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 datetime import datetime
|
|
from typing import Optional, Callable, Dict
|
|
|
|
from .config import (
|
|
UDP_HOST, UDP_PORT, UDP_BUFFER_SIZE,
|
|
SECRET_TOKEN, IMAGE_DIR,
|
|
VIDEO_FPS, VIDEO_CODEC
|
|
)
|
|
|
|
# Camera timeout - mark as inactive after this many seconds without frames
|
|
CAMERA_TIMEOUT_SECONDS = 5
|
|
|
|
|
|
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):
|
|
self.buffer = bytearray()
|
|
self.start_time = time.time()
|
|
self.receiving = True
|
|
|
|
def add_chunk(self, data: bytes) -> bool:
|
|
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]:
|
|
if not self.receiving or len(self.buffer) == 0:
|
|
return None
|
|
data = bytes(self.buffer)
|
|
self.reset()
|
|
return data
|
|
|
|
def reset(self):
|
|
self.buffer = bytearray()
|
|
self.start_time = None
|
|
self.receiving = False
|
|
|
|
|
|
class CameraRecorder:
|
|
"""Handles video recording for a single camera."""
|
|
|
|
def __init__(self, camera_id: str, output_dir: str):
|
|
self.camera_id = camera_id
|
|
self.output_dir = output_dir
|
|
self._writer: Optional[cv2.VideoWriter] = None
|
|
self._video_size: Optional[tuple] = None
|
|
self._recording = False
|
|
self._filename: Optional[str] = None
|
|
self._frame_count = 0
|
|
self._start_time: Optional[float] = None
|
|
|
|
@property
|
|
def is_recording(self) -> bool:
|
|
return self._recording
|
|
|
|
@property
|
|
def filename(self) -> Optional[str]:
|
|
return self._filename
|
|
|
|
@property
|
|
def duration(self) -> float:
|
|
if self._start_time:
|
|
return time.time() - self._start_time
|
|
return 0
|
|
|
|
@property
|
|
def frame_count(self) -> int:
|
|
return self._frame_count
|
|
|
|
def start(self) -> str:
|
|
if self._recording:
|
|
return self._filename
|
|
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
safe_id = self.camera_id.replace(":", "_").replace(".", "_")
|
|
self._filename = f"recording_{safe_id}_{timestamp}.avi"
|
|
self._recording = True
|
|
self._frame_count = 0
|
|
self._start_time = time.time()
|
|
return self._filename
|
|
|
|
def stop(self) -> dict:
|
|
if not self._recording:
|
|
return {"error": "Not recording"}
|
|
|
|
self._recording = False
|
|
result = {
|
|
"filename": self._filename,
|
|
"frames": self._frame_count,
|
|
"duration": self.duration
|
|
}
|
|
|
|
if self._writer:
|
|
self._writer.release()
|
|
self._writer = None
|
|
|
|
self._video_size = None
|
|
return result
|
|
|
|
def write_frame(self, frame: np.ndarray):
|
|
if not self._recording:
|
|
return
|
|
|
|
if self._writer is None:
|
|
self._video_size = (frame.shape[1], frame.shape[0])
|
|
fourcc = cv2.VideoWriter_fourcc(*VIDEO_CODEC)
|
|
video_path = os.path.join(self.output_dir, self._filename)
|
|
self._writer = cv2.VideoWriter(
|
|
video_path, fourcc, VIDEO_FPS, self._video_size
|
|
)
|
|
|
|
if self._writer and self._writer.isOpened():
|
|
self._writer.write(frame)
|
|
self._frame_count += 1
|
|
|
|
|
|
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,
|
|
device_registry=None):
|
|
self.host = host
|
|
self.port = port
|
|
self.image_dir = image_dir
|
|
self.on_frame = on_frame
|
|
self.device_registry = device_registry
|
|
|
|
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] = {}
|
|
|
|
# Per-camera recorders (keyed by device_id)
|
|
self._recorders: Dict[str, CameraRecorder] = {}
|
|
self._recordings_dir = os.path.join(os.path.dirname(image_dir), "recordings")
|
|
|
|
# IP to device_id mapping cache
|
|
self._ip_to_device: Dict[str, str] = {}
|
|
|
|
# Statistics
|
|
self.frames_received = 0
|
|
self.invalid_tokens = 0
|
|
self.decode_errors = 0
|
|
self.packets_received = 0
|
|
|
|
# Active cameras tracking: {device_id: {"last_frame": timestamp, "active": bool}}
|
|
self._active_cameras: Dict[str, dict] = {}
|
|
|
|
os.makedirs(self.image_dir, exist_ok=True)
|
|
os.makedirs(self._recordings_dir, exist_ok=True)
|
|
|
|
def set_device_registry(self, registry):
|
|
"""Set device registry for IP to device_id lookup."""
|
|
self.device_registry = registry
|
|
|
|
@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 device IDs."""
|
|
return [cid for cid, info in self._active_cameras.items() if info.get("active", False)]
|
|
|
|
def _get_device_id_from_ip(self, ip: str) -> Optional[str]:
|
|
"""Look up device_id from IP address using device registry."""
|
|
# Check cache first
|
|
if ip in self._ip_to_device:
|
|
return self._ip_to_device[ip]
|
|
|
|
# Look up in device registry
|
|
if self.device_registry:
|
|
for device in self.device_registry.all():
|
|
if device.address and device.address[0] == ip:
|
|
self._ip_to_device[ip] = device.id
|
|
return device.id
|
|
|
|
return None
|
|
|
|
def start(self) -> bool:
|
|
if self.is_running:
|
|
return False
|
|
|
|
self._stop_event.clear()
|
|
self._thread = threading.Thread(target=self._receive_loop, daemon=True)
|
|
self._thread.start()
|
|
|
|
# Start timeout checker
|
|
self._timeout_thread = threading.Thread(target=self._timeout_checker, daemon=True)
|
|
self._timeout_thread.start()
|
|
|
|
return True
|
|
|
|
def stop(self):
|
|
self._stop_event.set()
|
|
|
|
if self._sock:
|
|
try:
|
|
self._sock.close()
|
|
except Exception:
|
|
pass
|
|
self._sock = None
|
|
|
|
for recorder in self._recorders.values():
|
|
if recorder.is_recording:
|
|
recorder.stop()
|
|
|
|
self._cleanup_frames()
|
|
self._active_cameras.clear()
|
|
self._assemblers.clear()
|
|
self._recorders.clear()
|
|
self._ip_to_device.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 _timeout_checker(self):
|
|
"""Check for camera timeouts and mark them as inactive."""
|
|
while not self._stop_event.is_set():
|
|
time.sleep(1)
|
|
now = time.time()
|
|
|
|
for camera_id, info in list(self._active_cameras.items()):
|
|
last_frame = info.get("last_frame", 0)
|
|
was_active = info.get("active", False)
|
|
|
|
if now - last_frame > CAMERA_TIMEOUT_SECONDS:
|
|
if was_active:
|
|
self._active_cameras[camera_id]["active"] = False
|
|
# Remove the frame file so frontend shows default image
|
|
self._remove_camera_frame(camera_id)
|
|
|
|
def _remove_camera_frame(self, camera_id: str):
|
|
"""Remove the frame file for a camera."""
|
|
try:
|
|
filepath = os.path.join(self.image_dir, f"{camera_id}.jpg")
|
|
if os.path.exists(filepath):
|
|
os.remove(filepath)
|
|
except Exception:
|
|
pass
|
|
|
|
def _get_assembler(self, addr: tuple) -> FrameAssembler:
|
|
key = f"{addr[0]}:{addr[1]}"
|
|
if key not in self._assemblers:
|
|
self._assemblers[key] = FrameAssembler()
|
|
return self._assemblers[key]
|
|
|
|
def _get_recorder(self, camera_id: str) -> CameraRecorder:
|
|
if camera_id not in self._recorders:
|
|
self._recorders[camera_id] = CameraRecorder(camera_id, self._recordings_dir)
|
|
return self._recorders[camera_id]
|
|
|
|
def _receive_loop(self):
|
|
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
|
|
|
|
if not data.startswith(SECRET_TOKEN):
|
|
self.invalid_tokens += 1
|
|
continue
|
|
|
|
payload = data[len(SECRET_TOKEN):]
|
|
assembler = self._get_assembler(addr)
|
|
|
|
# Try to get device_id from IP, fallback to IP if not found
|
|
ip = addr[0]
|
|
device_id = self._get_device_id_from_ip(ip)
|
|
if not device_id:
|
|
# Fallback: use IP (without port to avoid duplicates)
|
|
device_id = ip.replace(".", "_")
|
|
|
|
if payload == b"START":
|
|
assembler.start_frame()
|
|
continue
|
|
elif payload == b"END":
|
|
frame_data = assembler.finish_frame()
|
|
if frame_data:
|
|
self._process_complete_frame(device_id, frame_data, addr)
|
|
continue
|
|
else:
|
|
if not assembler.receiving:
|
|
frame = self._decode_frame(payload)
|
|
if frame is not None:
|
|
self._process_frame(device_id, frame, addr)
|
|
else:
|
|
self.decode_errors += 1
|
|
else:
|
|
assembler.add_chunk(payload)
|
|
|
|
if self._sock:
|
|
self._sock.close()
|
|
self._sock = None
|
|
|
|
print("[UDP] Receiver stopped")
|
|
|
|
def _process_complete_frame(self, camera_id: str, frame_data: bytes, addr: tuple):
|
|
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):
|
|
self.frames_received += 1
|
|
|
|
# Update camera tracking
|
|
self._active_cameras[camera_id] = {
|
|
"last_frame": time.time(),
|
|
"active": True,
|
|
"addr": addr
|
|
}
|
|
|
|
# Save frame
|
|
self._save_frame(camera_id, frame)
|
|
|
|
# Record if recording is active for this camera
|
|
recorder = self._get_recorder(camera_id)
|
|
if recorder.is_recording:
|
|
recorder.write_frame(frame)
|
|
|
|
if self.on_frame:
|
|
self.on_frame(camera_id, frame, addr)
|
|
|
|
def _decode_frame(self, data: bytes) -> Optional[np.ndarray]:
|
|
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):
|
|
try:
|
|
filepath = os.path.join(self.image_dir, f"{camera_id}.jpg")
|
|
cv2.imwrite(filepath, frame)
|
|
except Exception:
|
|
pass
|
|
|
|
# === Recording API ===
|
|
|
|
def start_recording(self, camera_id: str) -> dict:
|
|
if camera_id not in self._active_cameras or not self._active_cameras[camera_id].get("active"):
|
|
return {"error": f"Camera {camera_id} not active"}
|
|
|
|
recorder = self._get_recorder(camera_id)
|
|
if recorder.is_recording:
|
|
return {"error": "Already recording", "filename": recorder.filename}
|
|
|
|
filename = recorder.start()
|
|
return {"status": "recording", "filename": filename, "camera_id": camera_id}
|
|
|
|
def stop_recording(self, camera_id: str) -> dict:
|
|
if camera_id not in self._recorders:
|
|
return {"error": f"No recorder for {camera_id}"}
|
|
|
|
recorder = self._recorders[camera_id]
|
|
if not recorder.is_recording:
|
|
return {"error": "Not recording"}
|
|
|
|
result = recorder.stop()
|
|
result["camera_id"] = camera_id
|
|
result["path"] = os.path.join(self._recordings_dir, result["filename"])
|
|
return result
|
|
|
|
def get_recording_status(self, camera_id: str = None) -> dict:
|
|
if camera_id:
|
|
if camera_id not in self._recorders:
|
|
return {"camera_id": camera_id, "recording": False}
|
|
recorder = self._recorders[camera_id]
|
|
return {
|
|
"camera_id": camera_id,
|
|
"recording": recorder.is_recording,
|
|
"filename": recorder.filename,
|
|
"duration": recorder.duration,
|
|
"frames": recorder.frame_count
|
|
}
|
|
|
|
result = {}
|
|
for cid, info in self._active_cameras.items():
|
|
if info.get("active"):
|
|
recorder = self._get_recorder(cid)
|
|
result[cid] = {
|
|
"recording": recorder.is_recording,
|
|
"filename": recorder.filename if recorder.is_recording else None,
|
|
"duration": recorder.duration if recorder.is_recording else 0
|
|
}
|
|
return result
|
|
|
|
def list_recordings(self) -> list:
|
|
try:
|
|
files = []
|
|
for f in os.listdir(self._recordings_dir):
|
|
if f.endswith(".avi"):
|
|
path = os.path.join(self._recordings_dir, f)
|
|
stat = os.stat(path)
|
|
files.append({
|
|
"filename": f,
|
|
"size": stat.st_size,
|
|
"created": stat.st_mtime
|
|
})
|
|
return sorted(files, key=lambda x: x["created"], reverse=True)
|
|
except Exception:
|
|
return []
|
|
|
|
def get_stats(self) -> dict:
|
|
recording_count = sum(1 for r in self._recorders.values() if r.is_recording)
|
|
active_count = sum(1 for info in self._active_cameras.values() if info.get("active"))
|
|
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": active_count,
|
|
"active_recordings": recording_count
|
|
}
|