Generates fiducial markers to improve inside-out tracking for VR headsets. 100% vibe-coded by Claude.

generate_fiducial_markers.py

#!/usr/bin/env python3
"""
Fiducial Marker Generator for VR Headset Tracking

This script generates unique, asymmetric fiducial markers with balanced 
black/white pixel distribution and sufficient Hamming distance separation.
Each marker is output as a separate page in a PDF file.
"""

import sys
import math
import random
import numpy as np
from itertools import product
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Image
from reportlab.lib.units import inch
from PIL import Image as PILImage
import io
import argparse


def calculate_grid_size(num_markers):
    """
    Calculate optimal grid size based on number of markers needed.
    
    For N×N grid, we have 2^(N²) possible patterns.
    We need to account for filtering out symmetric patterns and
    ensuring sufficient Hamming distance.
    """
    if num_markers <= 16:
        return 4  # 4×4 = 16 bits, 2^16 = 65,536 patterns
    elif num_markers <= 64:
        return 5  # 5×5 = 25 bits, 2^25 = 33,554,432 patterns
    elif num_markers <= 256:
        return 6  # 6×6 = 36 bits, 2^36 = 68,719,476,736 patterns
    else:
        return 7  # 7×7 = 49 bits, should handle very large sets


def pattern_to_bits(pattern):
    """Convert 2D pattern to 1D bit array."""
    return pattern.flatten()


def hamming_distance(bits1, bits2):
    """Calculate Hamming distance between two bit patterns."""
    return np.sum(bits1 != bits2)


def is_symmetric(pattern):
    """Check if pattern is symmetric under rotation or reflection."""
    # Check 90, 180, 270 degree rotations
    for k in range(1, 4):
        if np.array_equal(pattern, np.rot90(pattern, k)):
            return True
    
    # Check horizontal and vertical reflections
    if np.array_equal(pattern, np.fliplr(pattern)):
        return True
    if np.array_equal(pattern, np.flipud(pattern)):
        return True
    
    # Check diagonal reflections
    if np.array_equal(pattern, pattern.T):
        return True
    if np.array_equal(pattern, np.rot90(pattern.T)):
        return True
    
    return False


def has_balanced_distribution(pattern, tolerance=0.3):
    """Check if pattern has roughly equal black and white pixels."""
    total_pixels = pattern.size
    black_pixels = np.sum(pattern == 0)
    white_pixels = total_pixels - black_pixels
    
    ratio = min(black_pixels, white_pixels) / max(black_pixels, white_pixels)
    return ratio >= (1 - tolerance)


def generate_candidate_patterns(grid_size, num_candidates):
    """Generate candidate patterns that meet basic criteria."""
    candidates = []
    attempts = 0
    max_attempts = num_candidates * 100  # Prevent infinite loops
    
    while len(candidates) < num_candidates and attempts < max_attempts:
        # Generate random binary pattern
        pattern = np.random.randint(0, 2, size=(grid_size, grid_size))
        
        # Check if pattern meets criteria
        if not is_symmetric(pattern) and has_balanced_distribution(pattern):
            candidates.append(pattern)
        
        attempts += 1
    
    return candidates


def select_diverse_markers(candidates, num_markers, min_hamming_distance=3):
    """Select markers with sufficient Hamming distance separation."""
    if len(candidates) < num_markers:
        raise ValueError(f"Not enough valid candidates. Generated {len(candidates)}, need {num_markers}")
    
    selected = []
    selected_indices = set()
    candidate_bits = [pattern_to_bits(pattern) for pattern in candidates]
    
    # Start with a random candidate
    selected.append(candidates[0])
    selected_indices.add(0)
    selected_bits = [candidate_bits[0]]
    
    # Greedily select candidates with maximum minimum distance
    for _ in range(num_markers - 1):
        best_candidate = None
        best_bits = None
        best_min_distance = 0
        best_index = -1
        
        for i, (candidate, bits) in enumerate(zip(candidates, candidate_bits)):
            if i in selected_indices:
                continue
            
            # Calculate minimum distance to all selected markers
            min_distance = min(hamming_distance(bits, selected_bit) 
                             for selected_bit in selected_bits)
            
            if min_distance > best_min_distance:
                best_min_distance = min_distance
                best_candidate = candidate
                best_bits = bits
                best_index = i
        
        if best_candidate is not None and best_min_distance >= min_hamming_distance:
            selected.append(best_candidate)
            selected_indices.add(best_index)
            selected_bits.append(best_bits)
        else:
            # If we can't find a candidate with sufficient distance, take the best available
            if best_candidate is not None:
                selected.append(best_candidate)
                selected_indices.add(best_index)
                selected_bits.append(best_bits)
            else:
                break
    
    return selected


def create_marker_image(pattern, scale=50):
    """Create PIL Image from binary pattern."""
    height, width = pattern.shape
    
    # Scale up the pattern
    scaled_pattern = np.repeat(np.repeat(pattern, scale, axis=0), scale, axis=1)
    
    # Convert to 8-bit grayscale (0 = black, 255 = white)
    image_array = (scaled_pattern * 255).astype(np.uint8)
    
    return PILImage.fromarray(image_array, mode='L')


def generate_pdf(markers, output_filename):
    """Generate PDF with one marker per page."""
    doc = SimpleDocTemplate(output_filename, pagesize=letter)
    story = []
    
    for i, marker in enumerate(markers):
        # Create PIL image
        pil_image = create_marker_image(marker, scale=50)
        
        # Convert to bytes for reportlab
        img_buffer = io.BytesIO()
        pil_image.save(img_buffer, format='PNG')
        img_buffer.seek(0)
        
        # Create reportlab image
        img = Image(img_buffer)
        
        # Scale to fit page (leave some margin)
        max_size = 6 * inch
        img.drawHeight = max_size
        img.drawWidth = max_size
        
        story.append(img)
        
        # Add page break except for last marker
        if i < len(markers) - 1:
            from reportlab.platypus import PageBreak
            story.append(PageBreak())
    
    doc.build(story)


def main():
    parser = argparse.ArgumentParser(description='Generate fiducial markers for VR tracking')
    parser.add_argument('num_markers', type=int, help='Number of markers to generate')
    parser.add_argument('-o', '--output', default='fiducial_markers.pdf', 
                       help='Output PDF filename (default: fiducial_markers.pdf)')
    parser.add_argument('--min-hamming', type=int, default=3,
                       help='Minimum Hamming distance between markers (default: 3)')
    
    args = parser.parse_args()
    
    if args.num_markers <= 0:
        print("Error: Number of markers must be positive")
        sys.exit(1)
    
    print(f"Generating {args.num_markers} fiducial markers...")
    
    # Step 1: Determine grid size
    grid_size = calculate_grid_size(args.num_markers)
    print(f"Using {grid_size}×{grid_size} grid")
    
    # Step 2: Generate candidate patterns
    # Generate more candidates than needed to ensure diversity
    num_candidates = min(args.num_markers * 10, 10000)
    print(f"Generating {num_candidates} candidate patterns...")
    candidates = generate_candidate_patterns(grid_size, num_candidates)
    
    if len(candidates) < args.num_markers:
        print(f"Warning: Only generated {len(candidates)} valid candidates, need {args.num_markers}")
        print("Consider reducing the number of markers or relaxing constraints")
    
    # Step 3: Select diverse markers
    print("Selecting markers with sufficient Hamming distance...")
    selected_markers = select_diverse_markers(candidates, args.num_markers, args.min_hamming)
    
    if len(selected_markers) < args.num_markers:
        print(f"Warning: Only selected {len(selected_markers)} markers with sufficient diversity")
    
    # Step 4: Generate PDF
    print(f"Generating PDF: {args.output}")
    generate_pdf(selected_markers, args.output)
    
    print(f"Successfully generated {len(selected_markers)} fiducial markers in {args.output}")
    
    # Print some statistics
    if len(selected_markers) > 1:
        all_bits = [pattern_to_bits(marker) for marker in selected_markers]
        min_distance = min(hamming_distance(all_bits[i], all_bits[j]) 
                          for i in range(len(all_bits)) 
                          for j in range(i+1, len(all_bits)))
        print(f"Minimum Hamming distance between markers: {min_distance}")


if __name__ == '__main__':
    main()