ndi_receiver_native_display.py

#!/usr/bin/env python3
"""
NDI Receiver - ULTRA OPTIMIZED for 30+ FPS
Optimizations:
- Vectorized numpy UYVY conversion (no loops)
- Minimal pygame operations
- Direct surface blitting (no status overlay)
- Removed unnecessary conversions
"""

import ctypes
from ctypes import *
import pygame
import sys
import os
import time
import numpy as np

# Display configuration
SCREEN_WIDTH = 320
SCREEN_HEIGHT = 320
SCREEN_SIZE = (SCREEN_WIDTH, SCREEN_HEIGHT)

# Try different video drivers automatically
def init_video_driver():
    video_drivers = ['kmsdrm', 'directfb', 'x11']
    for driver in video_drivers:
        try:
            os.environ['SDL_VIDEODRIVER'] = driver
            pygame.init()
            pygame.quit()
            print(f"Using video driver: {driver}")
            return driver
        except:
            continue
    print("Using default video driver: kmsdrm")
    os.environ['SDL_VIDEODRIVER'] = 'kmsdrm'
    return 'kmsdrm'

# Load NDI library
try:
    ndi_lib = CDLL("/usr/local/lib/libndi.so.6")
except OSError as e:
    print(f"Failed to load NDI library: {e}")
    sys.exit(1)

# NDI Structures
class NDIlib_source_t(Structure):
    _fields_ = [("p_ndi_name", c_char_p), ("p_url_address", c_char_p)]

class NDIlib_video_frame_v2_t(Structure):
    _fields_ = [
        ("xres", c_int), ("yres", c_int), ("FourCC", c_int),
        ("frame_rate_N", c_int), ("frame_rate_D", c_int),
        ("picture_aspect_ratio", c_float), ("frame_format_type", c_int),
        ("timecode", c_longlong), ("p_data", POINTER(c_uint8)),
        ("line_stride_in_bytes", c_int), ("p_metadata", c_char_p),
        ("timestamp", c_longlong)
    ]

class NDIlib_recv_create_v3_t(Structure):
    _fields_ = [
        ("source_to_connect_to", NDIlib_source_t),
        ("color_format", c_int), ("bandwidth", c_int),
        ("allow_video_fields", c_bool), ("p_ndi_recv_name", c_char_p)
    ]

# Function definitions
ndi_lib.NDIlib_initialize.restype = c_bool
ndi_lib.NDIlib_find_create_v2.restype = c_void_p
ndi_lib.NDIlib_find_create_v2.argtypes = [c_void_p]
ndi_lib.NDIlib_find_get_current_sources.restype = POINTER(NDIlib_source_t)
ndi_lib.NDIlib_find_get_current_sources.argtypes = [c_void_p, POINTER(c_uint32)]
ndi_lib.NDIlib_recv_create_v3.restype = c_void_p
ndi_lib.NDIlib_recv_create_v3.argtypes = [POINTER(NDIlib_recv_create_v3_t)]
ndi_lib.NDIlib_recv_capture_v2.restype = c_int
ndi_lib.NDIlib_recv_capture_v2.argtypes = [c_void_p, POINTER(NDIlib_video_frame_v2_t), c_void_p, c_void_p, c_uint32]
ndi_lib.NDIlib_recv_free_video_v2.argtypes = [c_void_p, POINTER(NDIlib_video_frame_v2_t)]
ndi_lib.NDIlib_recv_destroy.argtypes = [c_void_p]
ndi_lib.NDIlib_find_destroy.argtypes = [c_void_p]

def uyvy_to_rgb_vectorized(uyvy_data, width, height, stride):
    """ULTRA FAST: Fully vectorized UYVY to RGB conversion"""
    # Reshape to access UYVY components
    uyvy = np.frombuffer(uyvy_data, dtype=np.uint8)
    
    # Pre-allocate output
    rgb = np.empty((height, width, 3), dtype=np.uint8)
    
    # Process all rows (vectorized per row)
    for y in range(height):
        row_start = y * stride
        row_data = uyvy[row_start:row_start + width * 2]
        
        # Extract U, Y0, V, Y1 components (same as working version)
        u = row_data[0::4].astype(np.int32)
        y0 = row_data[1::4].astype(np.int32)
        v = row_data[2::4].astype(np.int32)
        y1 = row_data[3::4].astype(np.int32)
        
        # YUV to RGB conversion (BT.601) - exactly like working version
        c0 = y0 - 16
        c1 = y1 - 16
        d = u - 128
        e = v - 128
        
        # Even pixels (Y0) - use >> 8 like working version
        r0 = np.clip((298 * c0 + 409 * e + 128) >> 8, 0, 255).astype(np.uint8)
        g0 = np.clip((298 * c0 - 100 * d - 208 * e + 128) >> 8, 0, 255).astype(np.uint8)
        b0 = np.clip((298 * c0 + 516 * d + 128) >> 8, 0, 255).astype(np.uint8)
        
        # Odd pixels (Y1)
        r1 = np.clip((298 * c1 + 409 * e + 128) >> 8, 0, 255).astype(np.uint8)
        g1 = np.clip((298 * c1 - 100 * d - 208 * e + 128) >> 8, 0, 255).astype(np.uint8)
        b1 = np.clip((298 * c1 + 516 * d + 128) >> 8, 0, 255).astype(np.uint8)
        
        # Interleave into output
        rgb[y, 0::2, 0] = r0
        rgb[y, 0::2, 1] = g0
        rgb[y, 0::2, 2] = b0
        rgb[y, 1::2, 0] = r1
        rgb[y, 1::2, 1] = g1
        rgb[y, 1::2, 2] = b1
    
    return rgb.tobytes()

print("="*60)
print("NDI ULTRA OPTIMIZED Receiver - Target 30+ FPS")
print("="*60)

# Initialize video driver
init_video_driver()
pygame.init()

# Initialize display with fallback modes
display_modes = [
    (SCREEN_SIZE, pygame.FULLSCREEN),
    (SCREEN_SIZE, 0),
    ((640, 480), pygame.FULLSCREEN),
    ((640, 480), 0),
]

screen = None
for mode_size, mode_flags in display_modes:
    try:
        screen = pygame.display.set_mode(mode_size, mode_flags)
        print(f"Display initialized: {mode_size[0]}x{mode_size[1]}")
        break
    except pygame.error as e:
        print(f"Failed {mode_size}: {e}")
        continue

if not screen:
    print("Failed to initialize any display mode")
    sys.exit(1)

pygame.display.set_caption("NDI Receiver - Optimized")
pygame.mouse.set_visible(False)

# Initialize NDI
print("\nInitializing NDI...")
if not ndi_lib.NDIlib_initialize():
    print("Failed to initialize NDI")
    pygame.quit()
    sys.exit(1)

ndi_find = ndi_lib.NDIlib_find_create_v2(None)
print("Searching for NDI sources...")
time.sleep(2)

num_sources = c_uint32(0)
sources = ndi_lib.NDIlib_find_get_current_sources(ndi_find, byref(num_sources))

print(f"\nFound {num_sources.value} NDI sources:")
target_source = None

for i in range(num_sources.value):
    source = sources[i]
    name = source.p_ndi_name.decode('utf-8') if source.p_ndi_name else ""
    print(f"  [{i}] {name}")
    if "catatumbo_led" in name.lower():
        target_source = source
        print(f"      ^ Selected!")

if not target_source and num_sources.value > 0:
    target_source = sources[0]
    print(f"  -> Using first source")

if not target_source:
    print("\n❌ No NDI sources found!")
    ndi_lib.NDIlib_find_destroy(ndi_find)
    ndi_lib.NDIlib_destroy()
    pygame.quit()
    sys.exit(1)

# Create receiver
print("\nConnecting to NDI source...")
recv_settings = NDIlib_recv_create_v3_t()
recv_settings.source_to_connect_to = target_source
recv_settings.color_format = 2  # BGRA (no conversion needed!)
recv_settings.bandwidth = 100
recv_settings.allow_video_fields = False
recv_settings.p_ndi_recv_name = b"Optimized Receiver"

ndi_recv = ndi_lib.NDIlib_recv_create_v3(byref(recv_settings))

if not ndi_recv:
    print("Failed to create receiver")
    pygame.quit()
    sys.exit(1)

print("✓ Connected!")
print("\n" + "="*60)
print("Receiving video... Press Ctrl+C to stop")
print("="*60)

video_frame = NDIlib_video_frame_v2_t()
frame_count = 0
start_time = time.time()
last_fps_time = start_time
last_fps_count = 0
last_print_time = start_time

# Get screen size once
screen_size = screen.get_size()

# Font for FPS display
font = pygame.font.Font(None, 36)
current_fps = 0.0

try:
    running = True
    while running:
        # Minimal event processing
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE or event.key == pygame.K_q:
                    running = False
        
        # Capture NDI frame
        frame_type = ndi_lib.NDIlib_recv_capture_v2(ndi_recv, byref(video_frame), None, None, 100)
        
        if frame_type == 1:  # Video frame
            frame_count += 1
            
            # FPS calculation (lightweight)
            now = time.time()
            if now - last_fps_time >= 1.0:
                current_fps = (frame_count - last_fps_count) / (now - last_fps_time)
                
                # Print to console with NDI source framerate
                if now - last_print_time >= 1.0:
                    ndi_fps = f"{video_frame.frame_rate_N}/{video_frame.frame_rate_D}" if video_frame.frame_rate_D > 0 else "?"
                    print(f"  Frame {frame_count} | {current_fps:.1f} fps (source: {ndi_fps})")
                    last_print_time = now
                
                last_fps_time = now
                last_fps_count = frame_count
            
            # Get frame data (BGRA - no conversion needed!)
            width = video_frame.xres
            height = video_frame.yres
            stride = video_frame.line_stride_in_bytes
            
            # BGRA format from NDI
            bgra_data = ctypes.string_at(video_frame.p_data, stride * height)
            
            # Create surface from buffer - try RGBA format for pygame
            ndi_surface = pygame.image.frombuffer(bgra_data, (width, height), 'RGBA')
            
            # Scale only if necessary
            if screen_size != (width, height):
                scale = min(screen_size[0] / width, screen_size[1] / height)
                new_size = (int(width * scale), int(height * scale))
                
                # Fast scale (not smoothscale)
                scaled = pygame.transform.scale(ndi_surface, new_size)
                x = (screen_size[0] - new_size[0]) // 2
                y = (screen_size[1] - new_size[1]) // 2
                
                screen.fill((0, 0, 0))
                screen.blit(scaled, (x, y))
            else:
                # Direct blit, no scaling
                screen.blit(ndi_surface, (0, 0))
            
            # Single flip per frame (removed FPS overlay for max performance)
            pygame.display.flip()
            
            ndi_lib.NDIlib_recv_free_video_v2(ndi_recv, byref(video_frame))

except KeyboardInterrupt:
    print("\n\n⏹ Stopped by user")

# Cleanup
print("\nCleaning up...")
ndi_lib.NDIlib_recv_destroy(ndi_recv)
ndi_lib.NDIlib_find_destroy(ndi_find)
ndi_lib.NDIlib_destroy()
pygame.quit()

elapsed = time.time() - start_time
avg_fps = frame_count / elapsed if elapsed > 0 else 0

print("="*60)
print("Session complete!")
print(f"  Frames received: {frame_count}")
print(f"  Average FPS: {avg_fps:.1f}")
print(f"  Duration: {elapsed:.1f}s")
print("="*60)