now using np to process UDP. fill(0) is slow though

This commit is contained in:
BOTAlex 2025-03-07 16:59:12 +01:00
parent cef892d728
commit 3f763204b6
2 changed files with 80 additions and 48 deletions

View File

@ -1,20 +1,11 @@
import struct import struct
import numpy as np
import ctypes
class DataPacket: class PacketHeader:
def __init__(self, buffer, offset, size): def __init__(self, headerFormat, dataView, rawDataSize):
self.frame_num = None self.frame_num: int|None = None
self.id = None self.id: int|None = None
self.buf_size = None self.image_buf_size: int|None = None
self.data = None
if len(buffer) < offset + 12: # Ensure buffer is large enough for metadata self.frame_num, self.id, self.image_buf_size = struct.unpack_from(headerFormat, dataView, 0)
return
self.frame_num, self.id, self.buf_size = struct.unpack_from("iii", buffer, offset)
current_index = offset + 12
message_length = size - current_index
if message_length < 0 or current_index + message_length > len(buffer):
return
self.data = buffer[current_index:current_index + message_length]

View File

@ -1,3 +1,4 @@
import struct
import cv2 import cv2
import numpy as np import numpy as np
import queue import queue
@ -13,10 +14,17 @@ import sys
from Camera.CameraState import CameraState from Camera.CameraState import CameraState
from Camera.ICameraSource import ICameraSource from Camera.ICameraSource import ICameraSource
import socket import socket
from .DataPacket import DataPacket # Python imports stupid. missed a single "." from .DataPacket import PacketHeader # Python imports stupid. missed a single "."
WAIT_TIME = 0.1 WAIT_TIME = 0.1
BUFFER_SIZE = 1024
NUM_FRAGMENTS = 48
HEADER_FORMAT = "iii"
RED = "\033[91m"
GREEN = "\033[92m"
RESET = "\033[0m"
# I HATE PYTHON!!!! slow fucking... # I HATE PYTHON!!!! slow fucking...
# Never know when value passed by value or ref # Never know when value passed by value or ref
@ -41,9 +49,12 @@ class UDP_Camera(ICameraSource):
self.host = "0.0.0.0" self.host = "0.0.0.0"
self.port = 3333 self.port = 3333
self.num_fragments = 48
self.num_loaded = 0 self.num_loaded = 0
self.packets = [None] * self.num_fragments self.packets = [None] * NUM_FRAGMENTS
self.rawDataBuffer = np.zeros(BUFFER_SIZE, dtype=np.uint8)
self.rawFullDataBuffer = np.zeros(BUFFER_SIZE*NUM_FRAGMENTS, dtype=np.uint8)
self.imageBuffView = memoryview(self.rawFullDataBuffer)
self.currentFrameNum = 0
self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.sock.bind((self.host, self.port)) self.sock.bind((self.host, self.port))
@ -57,53 +68,83 @@ class UDP_Camera(ICameraSource):
return return
else:
# We don't have a capture source to try yet, wait for one to show up in the GUI. bufferView = memoryview(self.rawDataBuffer)
if self.cancellation_event.wait(WAIT_TIME): n, senderAddr = self.sock.recvfrom_into(bufferView)
self.camera_status = CameraState.DISCONNECTED self.handle_packet(n, senderAddr)
def saveDataToBuff(self, dataView: memoryview, packet: PacketHeader):
header_size = struct.calcsize(HEADER_FORMAT)
payload_start = header_size
payload_end = header_size + packet.image_buf_size
payload = dataView[payload_start:payload_end]
#time.sleep(0.05)
offset = self.packets[0].image_buf_size * packet.id # Adjust offset based on payload size
self.imageBuffView[offset: offset + packet.image_buf_size] = payload
# print(f"Packet {packet.id}: {offset}->{offset + packet.image_buf_size}")
def resetImageBuffer(self):
self.rawFullDataBuffer = np.zeros(BUFFER_SIZE*NUM_FRAGMENTS, dtype=np.uint8)
def handle_packet(self, dataSize, senderAddr):
bufferView = memoryview(self.rawDataBuffer)
packet = PacketHeader(HEADER_FORMAT, bufferView, dataSize)
if packet.id < 0 or packet.id >= NUM_FRAGMENTS:
return return
data, addr = self.sock.recvfrom(1024) # Send acknowledgment
self.handle_packet(data, addr) self.sock.sendto(f"{packet.id}:{packet.frame_num}:ACK".encode(), senderAddr)
# print(packet.id)
def handle_packet(self, data, addr): if self.num_loaded > 0 and packet.id != 0:
packet = DataPacket(data, 0, len(data)) self.sock.sendto(f"ERR".encode(), senderAddr)
if packet.data is None or packet.id < 0 or packet.id >= self.num_fragments:
return
if packet.id == 0: if packet.id == 0 or packet.frame_num != self.currentFrameNum:
self.num_loaded = 1 self.packets: list[PacketHeader|None] = [None] * NUM_FRAGMENTS # Reset packets
self.packets = [None] * self.num_fragments self.num_loaded = 0
self.packets[packet.id] = packet self.currentFrameNum = packet.frame_num
elif self.packets[packet.id] is None: self.rawFullDataBuffer.fill(0)
self.packets[packet.id] = packet # print("Reset frame capture")
if (self.packets[packet.id] is None
and self.packets[0] is not None
or packet.id == 0 ):
self.num_loaded += 1 self.num_loaded += 1
self.packets[packet.id] = packet
self.saveDataToBuff(bufferView, packet)
if self.num_loaded >= self.num_fragments: # print(f"Got packet id: {packet.id} (total: )")
# if (self.num_loaded > NUM_FRAGMENTS * 0.75):
# formatted_list = "[" + ", ".join(f"{RED}x{RESET}" if item is None else f"{GREEN}x{RESET}" for item in self.packets) + "]"
# print(formatted_list)
if self.num_loaded >= NUM_FRAGMENTS:
if self.packets[0] is None: if self.packets[0] is None:
return return
current_frame = self.packets[0].frame_num current_frame = self.packets[0].frame_num
if any(p is None or p.frame_num != current_frame for p in self.packets): if any(p is None or p.frame_num != current_frame for p in self.packets):
print("Packet loss detected.") print("Packet loss detected.")
pass
else: else:
print("Whole data received. Processing...") #print("Whole data received. Processing...")
full_buffer = b''.join(p.data for p in self.packets) self.process_and_push_image()
np_buffer = np.frombuffer(full_buffer, dtype=np.uint8)
self.process_and_push_image(np_buffer)
self.num_loaded = 0 self.num_loaded = 0
self.packets: list[PacketHeader|None] = [None] * NUM_FRAGMENTS # Reset packets
# Send acknowledgment
self.sock.sendto(f"{packet.id}:ACK".encode(), addr)
def process_and_push_image(self, buf: np.ndarray): def process_and_push_image(self):
image = cv2.imdecode(buf, cv2.IMREAD_UNCHANGED) image = cv2.imdecode(self.rawFullDataBuffer, cv2.IMREAD_UNCHANGED)
if image is None: if image is None:
print(f"{Fore.YELLOW}[WARN] Frame drop. Corrupted JPEG.{Fore.RESET}") print(f"{Fore.YELLOW}[WARN] Frame drop. Corrupted JPEG.{Fore.RESET}")
return return
self.camera_status = CameraState.CONNECTED
current_frame_time = time.time() current_frame_time = time.time()
delta_time = current_frame_time - self.last_frame_time delta_time = current_frame_time - self.last_frame_time
self.last_frame_time = current_frame_time self.last_frame_time = current_frame_time