path_label.py

"""
Filename:   path_label.py
  Author:   Bennett Daniel

Path Labeler Tool
================

A GUI tool for labeling start/end points on costmap images to generate training data.

Features:
- Load and display costmap images
- Click to mark start (green) and end (red) points
- Save point coordinates to CSV
- Navigate between images with buttons

How to use:
1. Set paths in main()
2. Run: python path_label.py
3. Click to mark points:
   - First click: Start point
   - Second click: End point
   - Third click: Reset points
4. Use buttons to navigate images
5. Close window when done

Output CSV columns:
- costmap: Image filename
- start_x/y: Start coordinates 
- end_x/y: End coordinates
"""

import os
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.widgets import Button

class CostmapAnnotator:
    """A graphical tool for annotating start and end points on costmap images"""
    
    def __init__(self, costmap_dir, output_file='dataset.csv'):
        self.costmap_dir = costmap_dir
        self.output_file = output_file
        self.costmap_files = sorted([f for f in os.listdir(costmap_dir) if f.startswith('map_') and f.endswith('.png')])
        self.current_index = 0
        self.start_point = None
        self.end_point = None
        self.annotations = self.load_existing_annotations()
        
        # Setup the figure and axes
        self.fig, self.ax = plt.subplots(figsize=(10, 8))
        plt.subplots_adjust(bottom=0.2)
        
        # Add navigation buttons
        self.ax_prev = plt.axes([0.1, 0.05, 0.1, 0.075])
        self.ax_next = plt.axes([0.21, 0.05, 0.1, 0.075])
        self.ax_save = plt.axes([0.32, 0.05, 0.1, 0.075])
        self.ax_skip = plt.axes([0.43, 0.05, 0.1, 0.075])
        
        self.btn_prev = Button(self.ax_prev, 'Previous')
        self.btn_next = Button(self.ax_next, 'Next')
        self.btn_save = Button(self.ax_save, 'Save')
        self.btn_skip = Button(self.ax_skip, 'Skip')
        
        self.btn_prev.on_clicked(self.prev_image)
        self.btn_next.on_clicked(self.next_image)
        self.btn_save.on_clicked(self.save_annotations)
        self.btn_skip.on_clicked(self.skip_image)
        
        # Connect the click event
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        
        # Display the first image
        self.display_current_image()
        
    def load_existing_annotations(self):
        """Load existing annotations if available"""

        if os.path.exists(self.output_file):
            try:
                return pd.read_csv(self.output_file)
            except pd.errors.EmptyDataError:
                # File exists but is empty, create a new DataFrame
                print(f"Warning: {self.output_file} exists but is empty. Creating new DataFrame.")
                return pd.DataFrame(columns=['costmap', 'start_x', 'start_y', 'end_x', 'end_y'])
        else:
            # Create an empty DataFrame with the required columns
            return pd.DataFrame(columns=['costmap', 'start_x', 'start_y', 'end_x', 'end_y'])
    
    def display_current_image(self):
        """Display the current costmap image"""

        if self.current_index >= len(self.costmap_files):
            plt.close()
            print("All costmaps have been annotated!")
            return
            
        costmap_file = self.costmap_files[self.current_index]
        img_path = os.path.join(self.costmap_dir, costmap_file)
        
        # Load and display the image
        self.image = cv2.imread(img_path)
        self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        
        self.ax.clear()
        self.ax.imshow(self.image)
        self.ax.set_title(f'Costmap {self.current_index+1}/{len(self.costmap_files)}: {costmap_file}')
        
        # Check if this image already has annotations
        if not self.annotations.empty and 'costmap' in self.annotations.columns:
            existing = self.annotations[self.annotations['costmap'] == costmap_file]
            if not existing.empty:
                row = existing.iloc[0]
                self.start_point = (row['start_x'], row['start_y'])
                self.end_point = (row['end_x'], row['end_y'])
                
                # Display existing points
                if self.start_point:
                    self.ax.plot(self.start_point[0], self.start_point[1], 'go', markersize=10, label='Start')
                if self.end_point:
                    self.ax.plot(self.end_point[0], self.end_point[1], 'ro', markersize=10, label='End')
                
                if self.start_point and self.end_point:
                    self.ax.legend()
            else:
                self.start_point = None
                self.end_point = None
        else:
            self.start_point = None
            self.end_point = None
        
        self.fig.canvas.draw_idle()
    
    def on_click(self, event):
        """Handle mouse clicks on the image"""

        if event.inaxes != self.ax:
            return
            
        x, y = int(event.xdata), int(event.ydata)
        
        if self.start_point is None:
            # First click sets the start point
            self.start_point = (x, y)
            self.ax.plot(x, y, 'go', markersize=10, label='Start')
        elif self.end_point is None:
            # Second click sets the end point
            self.end_point = (x, y)
            self.ax.plot(x, y, 'ro', markersize=10, label='End')
            self.ax.legend()
            
            # Save the annotation for the current image
            self.save_current_annotation()
        else:
            # If both points are already set, reset and set start point
            self.ax.clear()
            self.ax.imshow(self.image)
            self.start_point = (x, y)
            self.end_point = None
            self.ax.plot(x, y, 'go', markersize=10, label='Start')
            
        self.fig.canvas.draw_idle()
    
    def save_current_annotation(self):
        """Save annotation for the current image"""
        
        costmap_file = self.costmap_files[self.current_index]
        
        # Remove existing annotation for this costmap if it exists
        if not self.annotations.empty and 'costmap' in self.annotations.columns:
            self.annotations = self.annotations[self.annotations['costmap'] != costmap_file]
        
        # Add the new annotation
        new_row = pd.DataFrame({
            'costmap': [costmap_file],
            'start_x': [self.start_point[0]],
            'start_y': [self.start_point[1]],
            'end_x': [self.end_point[0]],
            'end_y': [self.end_point[1]]
        })
        
        self.annotations = pd.concat([self.annotations, new_row], ignore_index=True)
        
        # Save to file
        self.annotations.to_csv(self.output_file, index=False)
        print(f"Saved annotation for {costmap_file}")
        
        # Automatically move to the next image
        self.next_image(None)
    
    def prev_image(self, event):
        """Go to the previous image"""

        if self.current_index > 0:
            self.current_index -= 1
            self.display_current_image()
    
    def next_image(self, event):
        """Go to the next image"""

        if self.current_index < len(self.costmap_files) - 1:
            self.current_index += 1
            self.display_current_image()
    
    def skip_image(self, event):
        """Skip the current image without saving annotations"""

        if self.current_index < len(self.costmap_files) - 1:
            self.current_index += 1
            self.display_current_image()
    
    def save_annotations(self, event):
        """Save all annotations to CSV file"""

        self.annotations.to_csv(self.output_file, index=False)
        print(f"All annotations saved to {self.output_file}")

def main():
    """Main function to run the Costmap Annotation Tool
    
    Configure the costmap directory and output file path here before running.
    """
    costmap_dir = '/Users/bennett/Documents/GitHub/path_planner_visualization/costmaps'
    output_file = '/Users/bennett/Documents/GitHub/path_planner_visualization/dataset.csv'
    
    print("Starting Costmap Annotation Tool...")
    print("Instructions:")
    print("1. First click: Set start point (green)")
    print("2. Second click: Set end point (red)")
    print("3. Third click: Reset points and set new start point")
    print("4. Use buttons to navigate between images")
    print("5. Close the window when finished")
    
    annotator = CostmapAnnotator(costmap_dir, output_file)
    plt.show()
    
    print(f"Annotation complete. Data saved to {output_file}")

if __name__ == "__main__":
    main()