createcurbs.py

import numpy as np
import json
import matplotlib.pyplot as plt
from shapely.geometry import LineString, Point
import math
import copy

def load_geojson(file_path):
    """Load GeoJSON file from path"""
    with open(file_path, 'r') as f:
        return json.load(f)

def create_spatial_index(features):
    """Create a spatial index from line features for efficient lookup"""
    # Filter for LineString features only
    line_features = [f for f in features if f['geometry']['type'] == 'LineString']
    line_index = {}
    for idx, feature in enumerate(line_features):
        coords = feature['geometry']['coordinates']
        start_point = tuple(coords[0])
        end_point = tuple(coords[-1])
        
        if start_point not in line_index:
            line_index[start_point] = []
        if end_point not in line_index:
            line_index[end_point] = []
            
        line_index[start_point].append(idx)
        line_index[end_point].append(idx)
    
    return line_features, line_index

def find_blocks(line_features, line_index):
    """Find blocks by detecting connected components in the line segment graph"""
    visited = set()  # Track visited line segments
    blocks = []      # Store detected blocks
    
    for idx, feature in enumerate(line_features):
        if idx in visited:
            continue
            
        block_segments = []
        queue = [idx]
        block_lines = set()
        
        while queue:
            line_idx = queue.pop(0)
            if line_idx in block_lines:
                continue
                
            block_lines.add(line_idx)
            visited.add(line_idx)
            
            coords = line_features[line_idx]['geometry']['coordinates']
            start = coords[0]
            end = coords[-1]
            dx = abs(end[0] - start[0])
            dy = abs(end[1] - start[1])
            is_horizontal = dx > dy
            
            block_segments.append({
                'index': line_idx,
                'coords': coords,
                'is_horizontal': is_horizontal
            })
            
            start_point = tuple(coords[0])
            end_point = tuple(coords[-1])
            
            connected_lines = set()
            for point in [start_point, end_point]:
                if point in line_index:
                    connected_lines.update(line_index[point])
            
            for connected_idx in connected_lines:
                if connected_idx not in block_lines:
                    queue.append(connected_idx)
        
        if len(block_segments) >= 3:
            blocks.append(block_segments)
    
    return blocks

def calculate_block_geometry(block_segments):
    """Calculate block center and bounding box"""
    all_coords = []
    for segment in block_segments:
        all_coords.extend(segment['coords'])
    
    # Calculate center
    center_x = sum(c[0] for c in all_coords) / len(all_coords)
    center_y = sum(c[1] for c in all_coords) / len(all_coords)
    center = (center_x, center_y)
    
    # Calculate bounding box
    min_x = min(c[0] for c in all_coords)
    max_x = max(c[0] for c in all_coords)
    min_y = min(c[1] for c in all_coords)
    max_y = max(c[1] for c in all_coords)
    bounds = (min_x, min_y, max_x, max_y)
    
    return center, bounds

def assign_cardinal_directions(block_segments):
    """Assign cardinal directions to block segments"""
    center, bounds = calculate_block_geometry(block_segments)
    min_x, min_y, max_x, max_y = bounds
    center_x, center_y = center
    
    for segment in block_segments:
        coords = segment['coords']
        seg_center_x = sum(c[0] for c in coords) / len(coords)
        seg_center_y = sum(c[1] for c in coords) / len(coords)
        
        if segment['is_horizontal']:
            # Horizontal segments are North or South
            if seg_center_y > center_y:
                segment['direction'] = 'N'  # North is above center
            else:
                segment['direction'] = 'S'  # South is below center
        else:
            # Vertical segments are East or West
            if seg_center_x > center_x:
                segment['direction'] = 'E'  # East is right of center
            else:
                segment['direction'] = 'W'  # West is left of center
    
    return block_segments, center

def resolve_direction_conflicts(block_segments):
    """Ensure each block has all four cardinal directions if possible"""
    direction_counts = {'N': 0, 'S': 0, 'E': 0, 'W': 0}
    direction_segments = {'N': [], 'S': [], 'E': [], 'W': []}
    
    for segment in block_segments:
        direction = segment['direction']
        direction_counts[direction] += 1
        direction_segments[direction].append(segment)
    
    def resolve_conflict(dir1, dir2):
        if direction_counts[dir1] > 1 and direction_counts[dir2] == 0:
            # Sort by length to find the shortest segment to reassign
            segments = direction_segments[dir1]
            segments.sort(key=lambda s: len(s['coords']), reverse=True)
            segment_to_move = segments.pop()
            
            # Update direction
            segment_to_move['direction'] = dir2
            direction_counts[dir1] -= 1
            direction_counts[dir2] += 1
            direction_segments[dir2].append(segment_to_move)
    
    resolve_conflict('N', 'S')
    resolve_conflict('S', 'N')
    resolve_conflict('E', 'W')
    resolve_conflict('W', 'E')
    
    return block_segments

def project_point_on_line(point, line_start, line_end):
    """Project a point onto a line segment"""
    px, py = point
    x1, y1 = line_start
    x2, y2 = line_end
    
    # Vector components
    dx = x2 - x1
    dy = y2 - y1
    
    # Squared length of the line segment
    length_squared = dx*dx + dy*dy
    if length_squared == 0:
        return line_start
    t = max(0, min(1, ((px - x1) * dx + (py - y1) * dy) / length_squared))
    projected_x = x1 + t * dx
    projected_y = y1 + t * dy
    
    return (projected_x, projected_y)

def distance(point1, point2):
    """Calculate Euclidean distance between two points"""
    return math.sqrt((point2[0] - point1[0])**2 + (point2[1] - point1[1])**2)

def split_segment_at_point(segment, point):
    """Split a segment at a specific point"""
    coords = segment['coords']
    
    min_distance = float('inf')
    insert_index = 0
    
    for i in range(len(coords) - 1):
        projected = project_point_on_line(point, coords[i], coords[i+1])
        dist = distance(point, projected)
        
        if dist < min_distance:
            min_distance = dist
            insert_index = i + 1
    
    before_coords = coords[:insert_index]
    before_coords.append(point)
    
    after_coords = [point] + coords[insert_index:]
    
    # Create the new segments with the same properties
    before = copy.deepcopy(segment)
    before['coords'] = before_coords
    
    after = copy.deepcopy(segment)
    after['coords'] = after_coords
    after['id'] = str(segment['id']) + '_split'
    
    return before, after

def snap_poles_to_segments(segments, pole_points, threshold=0.001):
    """Snap poles to segments and split segments at pole points"""
    working_segments = copy.deepcopy(segments)
    snapped_poles = []
    
    for pole_point in pole_points:
        # Find closest segment and point on segment
        closest_segment = None
        closest_point = None
        min_distance = float('inf')
        segment_index = -1
        
        for idx, segment in enumerate(working_segments):
            coords = segment['coords']
            
            # Check each line segment
            for i in range(len(coords) - 1):
                # Project the pole point onto this line segment
                projected = project_point_on_line(pole_point, coords[i], coords[i+1])
                dist = distance(pole_point, projected)
                
                if dist < min_distance:
                    min_distance = dist
                    closest_segment = segment
                    closest_point = projected
                    segment_index = idx
        
        if closest_segment and closest_point and min_distance < threshold:
            # Record the snapped pole
            snapped_poles.append({
                'coords': closest_point,
                'segment_label': closest_segment.get('label')
            })
            
            # Split the segment at the pole
            before, after = split_segment_at_point(closest_segment, closest_point)
            
            # Replace the original segment with the before part
            working_segments[segment_index] = before
            
            # Add the after part as a new segment
            working_segments.append(after)
    
    return working_segments, snapped_poles

def renumber_segments(segments):
    """Number segments sequentially within each block-direction group"""
    # Group segments by block and direction
    groups = {}
    for segment in segments:
        key = f"{segment['block_id']}{segment['direction']}"
        if key not in groups:
            groups[key] = []
        groups[key].append(segment)
    
    # For each group, sort and renumber
    renumbered_segments = []
    
    for key, group in groups.items():
        # Sort segments by position
        group.sort(key=lambda s: calculate_segment_position(s, key[-1]))
        
        # Number segments in order
        for idx, segment in enumerate(group):
            renumbered_segment = copy.deepcopy(segment)
            renumbered_segment['label'] = f"{segment['block_id']}{segment['direction']}{idx+1}"
            renumbered_segments.append(renumbered_segment)
    
    return renumbered_segments

def calculate_segment_position(segment, direction):
    """Calculate a position value for sorting segments within a direction"""
    # Calculate segment center
    coords = segment['coords']
    center_x = sum(c[0] for c in coords) / len(coords)
    center_y = sum(c[1] for c in coords) / len(coords)
    
    # For North/South segments, sort by X coordinate (west to east)
    if direction in ['N', 'S']:
        return center_x
    # For East/West segments, sort by Y coordinate (north to south)
    else:  # direction in ['E', 'W']
        return center_y

def process_curbs_and_poles(curb_geojson, pole_geojson=None):
    """Process curb and pole data to produce labeled segments"""
    # Extract features
    curb_features = curb_geojson['features']
    
    # 1. Create spatial index and identify line features
    line_features, line_index = create_spatial_index(curb_features)
    
    # 2. Find blocks
    blocks_raw = find_blocks(line_features, line_index)
    
    # 3. Process each block
    all_segments = []
    labeled_blocks = []
    
    for block_idx, block_segments in enumerate(blocks_raw):
        # 3.1 Assign cardinal directions
        directed_segments, block_center = assign_cardinal_directions(block_segments)
        
        # 3.2 Resolve direction conflicts
        resolved_segments = resolve_direction_conflicts(directed_segments)
        
        # 3.3 Prepare segments with block ID
        block_id = block_idx + 1  # 1-based indexing
        block_processed_segments = []
        
        for segment in resolved_segments:
            segment['block_id'] = block_id
            segment['id'] = segment['index']
            block_processed_segments.append(segment)
        
        labeled_blocks.append({
            'id': block_id,
            'center': block_center,
            'segments': block_processed_segments
        })
        
        all_segments.extend(block_processed_segments)
    
    # 4. Process poles if provided
    if pole_geojson:
        pole_features = pole_geojson['features']
        pole_points = [f['geometry']['coordinates'] for f in pole_features 
                      if f['geometry']['type'] == 'Point']
        
        # 4.1 Snap poles to segments
        split_segments, snapped_poles = snap_poles_to_segments(all_segments, pole_points)
        
        # 4.2 Renumber segments
        final_segments = renumber_segments(split_segments)
    else:
        # If no poles, just number the segments
        final_segments = renumber_segments(all_segments)
    
    return {
        'blocks': labeled_blocks,
        'segments': final_segments
    }

def create_output_geojson(processed_data, include_labels=True):
    """Create GeoJSON output from processed segments"""
    segments = processed_data['segments']
    
    features = []
    for segment in segments:
        feature = {
            'type': 'Feature',
            'geometry': {
                'type': 'LineString',
                'coordinates': segment['coords']
            },
            'properties': {}
        }
        
        # Add properties if labels are included
        if include_labels:
            feature['properties'] = {
                'label': segment['label'],
                'blockId': segment['block_id'],
                'direction': segment['direction'],
                'cardinalNumber': segment['label'][len(str(segment['block_id'])):]  # Extract the number part
            }
        
        features.append(feature)
    
    # Create the full GeoJSON
    geojson = {
        'type': 'FeatureCollection',
        'features': features
    }
    
    return geojson

def save_geojson(geojson, file_path):
    """Save GeoJSON to file"""
    with open(file_path, 'w') as f:
        json.dump(geojson, f, indent=2)

def visualize_blocks(processed_data, figsize=(12, 12), show_labels=True):
    """Visualize the detected blocks and segments"""
    segments = processed_data['segments']
    blocks = processed_data['blocks']
    
    # Create a figure
    plt.figure(figsize=figsize)
    
    # Direction colors - matching the HTML visualization
    direction_colors = {
        'N': '#FF5733',  # Red-orange
        'S': '#33FF57',  # Green
        'E': '#3357FF',  # Blue
        'W': '#F3FF33'   # Yellow
    }
    
    # Draw segments
    for segment in segments:
        coords = segment['coords']
        xs = [c[0] for c in coords]
        ys = [c[1] for c in coords]
        
        direction = segment['direction']
        plt.plot(xs, ys, color=direction_colors[direction], linewidth=2)
        
        # Optionally show labels at the midpoint of each segment
        if show_labels:
            mid_idx = len(coords) // 2
            mid_point = coords[mid_idx]
            plt.text(mid_point[0], mid_point[1], segment['label'], 
                     fontsize=8, ha='center', va='center',
                     bbox=dict(facecolor='white', alpha=0.7, edgecolor='none'))
    
    # Draw block numbers
    for block in blocks:
        center = block['center']
        plt.plot(center[0], center[1], 'ko', markersize=10)
        plt.text(center[0], center[1], str(block['id']), 
                 color='white', fontsize=8, ha='center', va='center')
    
    # Create a legend for directions
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor=direction_colors['N'], label='North'),
        Patch(facecolor=direction_colors['S'], label='South'),
        Patch(facecolor=direction_colors['E'], label='East'),
        Patch(facecolor=direction_colors['W'], label='West')
    ]
    plt.legend(handles=legend_elements, loc='upper right')
    
    plt.axis('equal')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.title('Detected Blocks and Curb Segments')
    plt.show()


if __name__ == "__main__":
    """Example usage of the curb segmentation algorithm"""
    
    # 1. Load the curb data
    curb_data = load_geojson('chinatown_curbs.geojson')
    
    # 2. Optionally load pole data
    try:
        pole_data = load_geojson('chinatown_poles.geojson')
        has_pole_data = True
    except FileNotFoundError:
        pole_data = None
        has_pole_data = False
        print("Pole data not found. Processing only with curb data.")
    
    # 3. Process the data
    if has_pole_data:
        processed_data = process_curbs_and_poles(curb_data, pole_data)
        print(f"Processed {len(processed_data['segments'])} segments with poles.")
    else:
        processed_data = process_curbs_and_poles(curb_data)
        print(f"Processed {len(processed_data['segments'])} segments without poles.")
        
    # 5. Create and save the output GeoJSON
    labeled_geojson = create_output_geojson(processed_data, include_labels=True)
    unlabeled_geojson = create_output_geojson(processed_data, include_labels=False)
    
    save_geojson(labeled_geojson, 'labeled_curbs.geojson')
    save_geojson(unlabeled_geojson, 'curbs.geojson')