coralnet-toolbox 0.0.75__py2.py3-none-any.whl → 0.0.76__py2.py3-none-any.whl

coralnet_toolbox/IO/QtExportMaskAnnotations.py

@@ -1,20 +1,19 @@
 import warnings
-warnings.filterwarnings("ignore", category=DeprecationWarning)
-
 import os
 import ujson as json
 
 import cv2
 import numpy as np
 import rasterio
-from PIL import Image
+from PIL import Image, ImageColor
 
-from PyQt5.QtCore import Qt
+from PyQt5.QtCore import Qt, pyqtSignal
+from PyQt5.QtGui import QColor, QPainter, QPen
 from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QGroupBox, QFormLayout,
                              QCheckBox, QComboBox, QLineEdit, QPushButton, QFileDialog,
-                             QApplication, QMessageBox, QLabel, QProgressDialog,
-                             QGroupBox, QListWidget, QAbstractItemView, QListWidgetItem,
-                             QButtonGroup, QScrollArea, QWidget)
+                             QApplication, QMessageBox, QLabel, QTableWidgetItem,
+                             QButtonGroup, QWidget, QTableWidget, QHeaderView,
+                             QAbstractItemView, QSpinBox, QRadioButton, QColorDialog)
 
 from coralnet_toolbox.Annotations.QtPatchAnnotation import PatchAnnotation
 from coralnet_toolbox.Annotations.QtPolygonAnnotation import PolygonAnnotation
@@ -22,12 +21,56 @@ from coralnet_toolbox.Annotations.QtRectangleAnnotation import RectangleAnnotati
 from coralnet_toolbox.Annotations.QtMultiPolygonAnnotation import MultiPolygonAnnotation
 
 from coralnet_toolbox.QtProgressBar import ProgressBar
-
 from coralnet_toolbox.Icons import get_icon
 
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
 
 # ----------------------------------------------------------------------------------------------------------------------
-# Classes
+# Helper Classes
+# ----------------------------------------------------------------------------------------------------------------------
+
+
+class ColorSwatchWidget(QWidget):
+    """A simple widget to display a color swatch with a border."""
+    def __init__(self, color, parent=None):
+        super().__init__(parent)
+        self.color = color
+        self.setFixedSize(24, 24)
+
+    def paintEvent(self, event):
+        painter = QPainter(self)
+        painter.setRenderHint(QPainter.Antialiasing)
+        
+        # Set the brush for the fill color
+        painter.setBrush(self.color)
+        
+        # Set the pen for the black border
+        pen = QPen(QColor("black"))
+        pen.setWidth(1)
+        painter.setPen(pen)
+        
+        # Draw the rectangle, adjusted inward so the border is fully visible
+        painter.drawRect(self.rect().adjusted(0, 0, -1, -1))
+
+    def setColor(self, color):
+        """Update the swatch's color and repaint."""
+        self.color = color
+        self.update()  # Triggers a repaint
+
+
+class ClickableColorSwatchWidget(ColorSwatchWidget):
+    """A ColorSwatchWidget that emits a clicked signal."""
+    clicked = pyqtSignal()
+
+    def mousePressEvent(self, event):
+        if event.button() == Qt.LeftButton:
+            self.clicked.emit()
+        super().mousePressEvent(event)
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+# Main Dialog Class
 # ----------------------------------------------------------------------------------------------------------------------
 
 
@@ -40,60 +83,70 @@ class ExportMaskAnnotations(QDialog):
         self.annotation_window = main_window.annotation_window
 
         self.setWindowIcon(get_icon("coral.png"))
-        self.setWindowTitle("Export Segmentation Masks")
-        self.resize(500, 250)
+        self.setWindowTitle("Export Annotations to Masks")
+        self.resize(1000, 800)
 
-        self.selected_labels = []
-        self.annotation_types = []
-        self.class_mapping = {}
-
-        # Create the layout
-        self.layout = QVBoxLayout(self)
-
-        # Setup the information layout
-        self.setup_info_layout()
-        # Setup the output directory and file format layout
-        self.setup_output_layout()
-        # Setup image selection layout
-        self.setup_image_selection_layout()
-        # Setup the annotation layout
-        self.setup_annotation_layout()
-        # Setup the label layout
-        self.setup_label_layout()
-        # Setup the mask format layout
-        self.setup_mask_format_layout()
-        # Setup the buttons layout
-        self.setup_buttons_layout()
+        self.mask_mode = 'semantic'  # 'semantic', 'sfm', or 'rgb'
+        self.rgb_background_color = QColor(0, 0, 0)
+
+        # Main layout for the dialog
+        self.main_layout = QVBoxLayout(self)
+
+        # Top section
+        top_section = QVBoxLayout()
+        self.setup_info_layout(parent_layout=top_section)
+        self.setup_output_layout(parent_layout=top_section)
+        self.setup_mask_format_layout(parent_layout=top_section)
+        self.main_layout.addLayout(top_section)
+
+        # Middle section
+        columns_layout = QHBoxLayout()
+        left_col = QVBoxLayout()
+        right_col = QVBoxLayout()
+
+        self.setup_annotation_layout(parent_layout=left_col)
+        self.setup_image_selection_layout(parent_layout=left_col)
+        self.setup_label_layout(parent_layout=right_col)
+
+        columns_layout.addLayout(left_col, 1)
+        columns_layout.addLayout(right_col, 2)
+        self.main_layout.addLayout(columns_layout)
+
+        # Bottom buttons
+        self.setup_buttons_layout(parent_layout=self.main_layout)
+
+        # Set initial mode and update UI
+        self.semantic_radio.setChecked(True)
+        self.update_ui_for_mode()
 
     def showEvent(self, event):
-        """Handle the show event"""
         super().showEvent(event)
-        # Update the labels in the label selection list
-        self.update_label_selection_list()
+        self.update_ui_for_mode()
 
-    def setup_info_layout(self):
-        """
-        Set up the layout and widgets for the info layout.
-        """
+    def setup_info_layout(self, parent_layout=None):
         group_box = QGroupBox("Information")
         layout = QVBoxLayout()
-
-        # Create a QLabel with explanatory text and hyperlink
-        info_label = QLabel("Export Annotations to Segmentation Masks")
-
-        info_label.setOpenExternalLinks(True)
+        info_text = (
+            "This tool exports annotations to image masks for three primary use cases:<br><br>"
+            "<b>1. Semantic Segmentation (Integer IDs):</b> Creates masks where each class (e.g., coral, rock) "
+            "is represented by a unique integer value (1, 2, 3...). The background is typically 0. These are used "
+            "to train machine learning models.<br><br>"
+            "<b>2. Structure from Motion (SfM) (Binary Mask):</b> Creates masks where a foreground value "
+            "(e.g., 255) represents objects to keep, and a background value (e.g., 0) represents areas to "
+            "ignore. This is used by software like Metashape to improve 3D model reconstruction.<br><br>"
+            "<b>3. Visualization (RGB Colors):</b> Creates a human-readable color mask using the colors "
+            "assigned to each label. Ideal for reports, presentations, and qualitative analysis."
+        )
+        info_label = QLabel(info_text)
         info_label.setWordWrap(True)
         layout.addWidget(info_label)
-
         group_box.setLayout(layout)
-        self.layout.addWidget(group_box)
+        parent_layout.addWidget(group_box)
 
-    def setup_output_layout(self):
-        """Setup the output directory and file format layout."""
+    def setup_output_layout(self, parent_layout=None):
         groupbox = QGroupBox("Output Directory and File Format")
         layout = QFormLayout()
 
-        # Output directory selection
         output_dir_layout = QHBoxLayout()
         self.output_dir_edit = QLineEdit()
         self.output_dir_button = QPushButton("Browse...")
@@ -102,18 +155,58 @@ class ExportMaskAnnotations(QDialog):
         output_dir_layout.addWidget(self.output_dir_button)
         layout.addRow("Output Directory:", output_dir_layout)
 
-        # Output folder name
-        self.output_name_edit = QLineEdit("")
+        self.output_name_edit = QLineEdit("masks")
         layout.addRow("Folder Name:", self.output_name_edit)
 
         groupbox.setLayout(layout)
-        self.layout.addWidget(groupbox)
+        parent_layout.addWidget(groupbox)
+
+    def setup_mask_format_layout(self, parent_layout=None):
+        groupbox = QGroupBox("Export Mode and Format")
+        main_layout = QVBoxLayout()
+
+        # Mode Selection
+        mode_layout = QHBoxLayout()
+        self.semantic_radio = QRadioButton("Semantic Segmentation (Integer IDs)")
+        self.sfm_radio = QRadioButton("Structure from Motion (Binary Mask)")
+        self.rgb_radio = QRadioButton("Visualization (RGB Colors)")
+
+        self.mode_group = QButtonGroup(self)
+        self.mode_group.addButton(self.semantic_radio)
+        self.mode_group.addButton(self.sfm_radio)
+        self.mode_group.addButton(self.rgb_radio)
+        self.mode_group.buttonClicked.connect(self.update_ui_for_mode)
+
+        mode_layout.addWidget(self.semantic_radio)
+        mode_layout.addWidget(self.sfm_radio)
+        mode_layout.addWidget(self.rgb_radio)
+        main_layout.addLayout(mode_layout)
+
+        # Format and Options
+        options_layout = QFormLayout()
+        
+        # General file format
+        self.file_format_combo = QComboBox()
+        self.file_format_combo.addItems([".png", ".bmp", ".tif"])
+        self.file_format_combo.currentTextChanged.connect(self.update_georef_availability)
+        options_layout.addRow("File Format:", self.file_format_combo)
 
-    def setup_image_selection_layout(self):
-        """Setup the image selection layout."""
+        # Georeferencing
+        self.preserve_georef_checkbox = QCheckBox("Preserve georeferencing (if available)")
+        self.preserve_georef_checkbox.setChecked(True)
+        options_layout.addRow(self.preserve_georef_checkbox)
+        self.georef_note = QLabel("Note: Georeferencing is only supported for TIF format.")
+        self.georef_note.setStyleSheet("color: #666; font-style: italic;")
+        options_layout.addRow(self.georef_note)
+
+        main_layout.addLayout(options_layout)
+        groupbox.setLayout(main_layout)
+        parent_layout.addWidget(groupbox)
+        self.update_georef_availability()
+
+    def setup_image_selection_layout(self, parent_layout=None):
         group_box = QGroupBox("Apply To")
         layout = QVBoxLayout()
-
         self.apply_filtered_checkbox = QCheckBox("▼ Apply to filtered images")
         self.apply_prev_checkbox = QCheckBox("↑ Apply to previous images")
         self.apply_next_checkbox = QCheckBox("↓ Apply to next images")
@@ -130,24 +223,18 @@ class ExportMaskAnnotations(QDialog):
         self.apply_group.addButton(self.apply_next_checkbox)
         self.apply_group.addButton(self.apply_all_checkbox)
         self.apply_group.setExclusive(True)
-
         group_box.setLayout(layout)
-        self.layout.addWidget(group_box)
+        parent_layout.addWidget(group_box)
 
-    def setup_annotation_layout(self):
-        """Setup the annotation types layout."""
+    def setup_annotation_layout(self, parent_layout=None):
         groupbox = QGroupBox("Annotations to Include")
         layout = QVBoxLayout()
-
-        # Annotation types checkboxes
         self.patch_checkbox = QCheckBox("Patch Annotations")
         self.patch_checkbox.setChecked(True)
         self.rectangle_checkbox = QCheckBox("Rectangle Annotations")
         self.rectangle_checkbox.setChecked(True)
         self.polygon_checkbox = QCheckBox("Polygon Annotations")
         self.polygon_checkbox.setChecked(True)
-
-        # Include negative samples
         self.include_negative_samples_checkbox = QCheckBox("Include negative samples")
         self.include_negative_samples_checkbox.setChecked(True)
 
@@ -155,438 +242,486 @@ class ExportMaskAnnotations(QDialog):
         layout.addWidget(self.rectangle_checkbox)
         layout.addWidget(self.polygon_checkbox)
         layout.addWidget(self.include_negative_samples_checkbox)
-
         groupbox.setLayout(layout)
-        self.layout.addWidget(groupbox)
+        parent_layout.addWidget(groupbox)
 
-    def setup_label_layout(self):
-        """Setup the label selection layout."""
-        groupbox = QGroupBox("Labels to Include")
+    def setup_label_layout(self, parent_layout=None):
+        groupbox = QGroupBox("Labels to Include / Rasterization Order")
         layout = QVBoxLayout()
-
-        # Label selection
-        self.label_selection_label = QLabel("Select Labels:")
-        layout.addWidget(self.label_selection_label)
-
-        # Create a scroll area for the labels
-        scroll_area = QScrollArea()
-        scroll_area.setWidgetResizable(True)
-
-        # Create a widget to hold the checkboxes
-        self.label_container = QWidget()
-        self.label_container.setMinimumHeight(200)  # Set a minimum height for the container
-        self.label_layout = QVBoxLayout(self.label_container)
-        self.label_layout.setSizeConstraint(QVBoxLayout.SetMinAndMaxSize)  # Respect widget sizes
-
-        scroll_area.setWidget(self.label_container)
-        layout.addWidget(scroll_area)
-
-        # Store the checkbox references
-        self.label_checkboxes = []
-
-        groupbox.setLayout(layout)
-        self.layout.addWidget(groupbox)
-
-    def setup_mask_format_layout(self):
-        """Setup the mask format layout."""
-        groupbox = QGroupBox("Mask Format")
-        layout = QFormLayout()
-
-        # File format combo box
-        self.file_format_combo = QComboBox()
-        self.file_format_combo.addItems([".png", ".bmp", ".tif"])
-        self.file_format_combo.setEditable(True)
-        layout.addRow("File Format:", self.file_format_combo)
-
-        # Add option to preserve georeferencing
-        self.preserve_georef_checkbox = QCheckBox("Preserve georeferencing (if available)")
-        self.preserve_georef_checkbox.setChecked(True)
-        layout.addRow("", self.preserve_georef_checkbox)
-
-        # Connect the checkbox to update based on file format
-        self.file_format_combo.currentTextChanged.connect(self.update_georef_availability)
-
-        # Add a note about georeferencing
-        self.georef_note = QLabel("Note: Georeferencing can only be preserved with TIF format")
-        self.georef_note.setStyleSheet("color: #666; font-style: italic;")
-        layout.addRow("", self.georef_note)
-
+        self.label_table = QTableWidget()
+        self.label_table.setColumnCount(3)
+        self.label_table.setSelectionBehavior(QAbstractItemView.SelectRows)
+        self.label_table.setSelectionMode(QAbstractItemView.SingleSelection)
+        header = self.label_table.horizontalHeader()
+        header.setSectionResizeMode(0, QHeaderView.ResizeToContents)
+        header.setSectionResizeMode(1, QHeaderView.Stretch)
+        header.setSectionResizeMode(2, QHeaderView.ResizeToContents)
+        layout.addWidget(self.label_table)
+        
+        button_layout = QHBoxLayout()
+        button_layout.addStretch(1)
+        self.move_up_button = QPushButton("▲ Move Up")
+        self.move_down_button = QPushButton("▼ Move Down")
+        self.move_up_button.clicked.connect(self.move_row_up)
+        self.move_down_button.clicked.connect(self.move_row_down)
+        button_layout.addWidget(self.move_up_button)
+        button_layout.addWidget(self.move_down_button)
+        button_layout.addStretch(1)
+        layout.addLayout(button_layout)
+        
+        order_note = QLabel(
+            "<b>Layer Order is Important:</b> Labels lower in the list will be drawn on top of labels "
+            "higher in the list. For overlapping annotations, only the topmost class will appear."
+        )
+        order_note.setStyleSheet("color: #666;")
+        order_note.setWordWrap(True)
+        layout.addWidget(order_note)
+        
         groupbox.setLayout(layout)
-        self.layout.addWidget(groupbox)
+        parent_layout.addWidget(groupbox)
 
-        # Initial update based on default format
-        self.update_georef_availability()
-
-    def setup_buttons_layout(self):
-        """Setup the buttons layout."""
+    def setup_buttons_layout(self, parent_layout=None):
         button_layout = QHBoxLayout()
-
+        button_layout.addStretch(1)
         self.export_button = QPushButton("Export")
-        self.export_button.clicked.connect(self.export_masks)
+        self.export_button.clicked.connect(self.run_export_process)
         self.cancel_button = QPushButton("Cancel")
         self.cancel_button.clicked.connect(self.reject)
-
         button_layout.addWidget(self.export_button)
         button_layout.addWidget(self.cancel_button)
-
-        self.layout.addLayout(button_layout)
-
+        parent_layout.addLayout(button_layout)
+
+    def update_ui_for_mode(self):
+        """Update the UI dynamically based on the selected export mode."""
+        if self.semantic_radio.isChecked():
+            self.mask_mode = 'semantic'
+        elif self.sfm_radio.isChecked():
+            self.mask_mode = 'sfm'
+        elif self.rgb_radio.isChecked():
+            self.mask_mode = 'rgb'
+        
+        self.populate_label_table()
+    
+    def populate_label_table(self):
+        """Populate the label table based on the current mode."""
+        self.label_table.blockSignals(True)
+        self.label_table.setRowCount(0)
+
+        # Set table headers based on mode
+        headers = ["Include", "Label Name"]
+        if self.mask_mode in ['semantic', 'sfm']:
+            headers.append("Mask Value")
+        elif self.mask_mode == 'rgb':
+            headers.append("Color Preview")
+        self.label_table.setHorizontalHeaderLabels(headers)
+
+        # --- BACKGROUND ROW (ROW 0) ---
+        self.label_table.insertRow(0)
+        checkbox_widget = self.create_centered_checkbox(checked=True)
+        self.label_table.setCellWidget(0, 0, checkbox_widget)
+        
+        label_item = QTableWidgetItem("Background")
+        label_item.setFlags(label_item.flags() & ~Qt.ItemIsEditable)
+        label_item.setData(Qt.UserRole, "Background")
+        self.label_table.setItem(0, 1, label_item)
+
+        if self.mask_mode in ['semantic', 'sfm']:
+            spinbox = QSpinBox()
+            spinbox.setRange(0, 255)
+            spinbox.setValue(0)
+            self.label_table.setCellWidget(0, 2, spinbox)
+        elif self.mask_mode == 'rgb':
+            swatch = ClickableColorSwatchWidget(self.rgb_background_color)
+            swatch.clicked.connect(self.pick_background_color)
+            # Create a container widget to center the swatch
+            container_widget = QWidget()
+            layout = QHBoxLayout(container_widget)
+            layout.addWidget(swatch)
+            layout.setAlignment(Qt.AlignCenter)
+            layout.setContentsMargins(0, 0, 0, 0)
+            self.label_table.setCellWidget(0, 2, container_widget)
+
+        # --- LABEL ROWS ---
+        for i, label in enumerate(self.label_window.labels):
+            row = i + 1
+            self.label_table.insertRow(row)
+
+            # Column 0: Include Checkbox
+            checkbox_widget = self.create_centered_checkbox(checked=True)
+            self.label_table.setCellWidget(row, 0, checkbox_widget)
+            
+            # Column 1: Label Name
+            label_item = QTableWidgetItem(label.short_label_code)
+            label_item.setFlags(label_item.flags() & ~Qt.ItemIsEditable)
+            label_item.setData(Qt.UserRole, label.short_label_code)
+            self.label_table.setItem(row, 1, label_item)
+
+            # Column 2: Mode-dependent widget
+            if self.mask_mode == 'semantic':
+                spinbox = QSpinBox()
+                spinbox.setRange(0, 255)
+                spinbox.setValue(i + 1)
+                self.label_table.setCellWidget(row, 2, spinbox)
+            elif self.mask_mode == 'sfm':
+                spinbox = QSpinBox()
+                spinbox.setRange(0, 255)
+                spinbox.setValue(255)  # Default foreground value
+                self.label_table.setCellWidget(row, 2, spinbox)
+            elif self.mask_mode == 'rgb':
+                try:
+                    q_color = QColor(label.color)
+                except Exception:
+                    q_color = QColor("#FFFFFF")  # Default to white on error
+                
+                swatch = ColorSwatchWidget(q_color)
+                cell_widget = QWidget()
+                layout = QHBoxLayout(cell_widget)
+                layout.addWidget(swatch)
+                layout.setAlignment(Qt.AlignCenter)
+                layout.setContentsMargins(0, 0, 0, 0)
+                self.label_table.setCellWidget(row, 2, cell_widget)
+        
+        if self.label_table.rowCount() > 0:
+            self.label_table.selectRow(0)
+        self.label_table.blockSignals(False)
+        
+    def pick_background_color(self):
+        color = QColorDialog.getColor(self.rgb_background_color, self, "Select Background Color")
+        if color.isValid():
+            self.rgb_background_color = color
+            swatch_container = self.label_table.cellWidget(0, 2)
+            if swatch_container:
+                # Find the swatch inside the container
+                swatch = swatch_container.findChild(ClickableColorSwatchWidget)
+                if swatch:
+                    swatch.setColor(color)
+
+    def create_centered_checkbox(self, checked=True):
+        checkbox = QCheckBox()
+        checkbox.setChecked(checked)
+        widget = QWidget()
+        layout = QHBoxLayout(widget)
+        layout.addWidget(checkbox)
+        layout.setAlignment(Qt.AlignCenter)
+        layout.setContentsMargins(0, 0, 0, 0)
+        return widget
+        
     def browse_output_dir(self):
-        """Open a file dialog to select the output directory."""
-        options = QFileDialog.Options()
-        directory = QFileDialog.getExistingDirectory(
-            self, "Select Output Directory", "", options=options
-        )
+        directory = QFileDialog.getExistingDirectory(self, "Select Output Directory")
         if directory:
             self.output_dir_edit.setText(directory)
 
-    def update_label_selection_list(self):
-        """Update the label selection list with labels from the label window."""
-        # Clear existing checkboxes
-        for checkbox in self.label_checkboxes:
-            self.label_layout.removeWidget(checkbox)
-            checkbox.deleteLater()
-        self.label_checkboxes = []
-
-        # Create a checkbox for each label
-        for label in self.label_window.labels:
-            checkbox = QCheckBox(label.short_label_code)
-            checkbox.setChecked(True)  # Default to checked
-            checkbox.setProperty("label", label)  # Store the label object
-            self.label_checkboxes.append(checkbox)
-            self.label_layout.addWidget(checkbox)
-
-    def update_georef_availability(self):
-        """Update georeferencing checkbox availability based on file format"""
-        current_format = self.file_format_combo.currentText().lower()
-        is_tif = '.tif' in current_format
-
-        self.preserve_georef_checkbox.setEnabled(is_tif)
-        if not is_tif:
-            self.preserve_georef_checkbox.setChecked(False)
-            self.georef_note.setStyleSheet("color: red; font-style: italic;")
-        else:
-            self.georef_note.setStyleSheet("color: #666; font-style: italic;")
-
     def get_selected_image_paths(self):
-        """
-        Get the selected image paths based on the options.
-
-        :return: List of selected image paths
-        """
-        # Current image path showing
         current_image_path = self.annotation_window.current_image_path
         if not current_image_path:
             return []
 
-        # Determine which images to export annotations for
         if self.apply_filtered_checkbox.isChecked():
             return self.image_window.table_model.filtered_paths
         elif self.apply_prev_checkbox.isChecked():
             if current_image_path in self.image_window.table_model.filtered_paths:
                 current_index = self.image_window.table_model.get_row_for_path(current_image_path)
                 return self.image_window.table_model.filtered_paths[:current_index + 1]
-            else:
-                return [current_image_path]
         elif self.apply_next_checkbox.isChecked():
             if current_image_path in self.image_window.table_model.filtered_paths:
                 current_index = self.image_window.table_model.get_row_for_path(current_image_path)
                 return self.image_window.table_model.filtered_paths[current_index:]
-            else:
-                return [current_image_path]
         elif self.apply_all_checkbox.isChecked():
             return self.image_window.raster_manager.image_paths
-        else:
-            # Only apply to the current image
-            return [current_image_path]
+        
+        return [current_image_path]
 
-    def export_class_mapping(self, output_path):
-        """Export the class mapping to a JSON file."""
-        mapping_file = os.path.join(output_path, "class_mapping.json")
-
-        with open(mapping_file, 'w') as f:
-            json.dump(self.class_mapping, f, indent=4)
-
-        if not os.path.exists(mapping_file):
-            print(f"Warning: Failed to save class mapping to {mapping_file}")
-
-    def export_masks(self):
-        """Export segmentation masks based on selected annotations and labels."""
-        # Validate inputs
+    def validate_inputs(self):
         if not self.output_dir_edit.text():
-            QMessageBox.warning(self,
-                                "Missing Output Directory",
+            QMessageBox.warning(self, 
+                                "Missing Input", 
                                 "Please select an output directory.")
-            return
-
-        # Check if at least one annotation type is selected
-        if not any([self.patch_checkbox.isChecked(),
-                    self.rectangle_checkbox.isChecked(),
+            return False
+        if not any([self.patch_checkbox.isChecked(), 
+                    self.rectangle_checkbox.isChecked(), 
                     self.polygon_checkbox.isChecked()]):
-            QMessageBox.warning(self,
-                                "No Annotation Type Selected",
+            QMessageBox.warning(self, 
+                                "Missing Input", 
                                 "Please select at least one annotation type.")
+            return False
+        return True
+
+    def run_export_process(self):
+        if not self.validate_inputs():
             return
 
-        # Check for checked items
-        self.selected_labels = []
-        for checkbox in self.label_checkboxes:
-            if checkbox.isChecked():
-                self.selected_labels.append(checkbox.property("label"))
-
-        # Check if at least one label is selected
-        if not self.selected_labels:
-            QMessageBox.warning(self,
-                                "No Labels Selected",
-                                "Please select at least one label.")
+        self.labels_to_render = []
+        self.background_value = 0 
+        
+        # --- Collect data from UI based on mode ---
+        if self.mask_mode in ['semantic', 'sfm']:
+            used_mask_values = {}
+            for i in range(self.label_table.rowCount()):
+                if not self.label_table.cellWidget(i, 0).findChild(QCheckBox).isChecked():
+                    continue
+                
+                label_code = self.label_table.item(i, 1).data(Qt.UserRole)
+                mask_value = self.label_table.cellWidget(i, 2).value()
+
+                if mask_value not in used_mask_values:
+                    used_mask_values[mask_value] = []
+                used_mask_values[mask_value].append(label_code)
+
+                if label_code == "Background":
+                    self.background_value = mask_value
+                else:
+                    label = next((l for l in self.label_window.labels if l.short_label_code == label_code), None)
+                    if label:
+                        self.labels_to_render.append((label, mask_value))
+            
+            # Check for duplicate values
+            duplicate_values = {v: l for v, l in used_mask_values.items() if len(l) > 1}
+            if duplicate_values:
+                msg = "Warning: The following mask values are used by multiple labels:\n" + \
+                      "\n".join([f"Value {v}: {', '.join(l)}" for v, l in duplicate_values.items()]) + \
+                      "\nThis may cause unexpected behavior. Continue?"
+                if QMessageBox.warning(self, 
+                                       "Duplicate Values", 
+                                       msg, 
+                                       QMessageBox.Yes | QMessageBox.No) == QMessageBox.No:
+                    return
+
+        elif self.mask_mode == 'rgb':
+            self.background_value = self.rgb_background_color.getRgb()[:3]  # (R, G, B) tuple
+            for i in range(1, self.label_table.rowCount()):  # Skip background
+                if self.label_table.cellWidget(i, 0).findChild(QCheckBox).isChecked():
+                    label_code = self.label_table.item(i, 1).data(Qt.UserRole)
+                    label = next((l for l in self.label_window.labels if l.short_label_code == label_code), None)
+                    if label:
+                        try:
+                            # Check if label.color is already a QColor object
+                            if isinstance(label.color, QColor):
+                                color_tuple = label.color.getRgb()[:3]  # Extract (R,G,B) and ignore alpha
+                            else:
+                                # Otherwise, assume it's a string (hex code) and convert it
+                                color_tuple = ImageColor.getrgb(label.color)
+                            self.labels_to_render.append((label, color_tuple))
+                        except (ValueError, TypeError) as e:
+                            print(f"Warning: Invalid color format for label "
+                                  f"'{label.short_label_code}': {label.color}. Error: {e}. Skipping.")
+
+        # --- Check if any labels are selected to be drawn ---
+        if not self.labels_to_render:
+            QMessageBox.warning(self, "No Labels Selected", "Please select at least one label to include in the masks.")
             return
 
+        # --- Setup paths and progress bar ---
         output_dir = self.output_dir_edit.text()
         folder_name = self.output_name_edit.text().strip()
-        file_format = self.file_format_combo.currentText()
-
-        # Ensure file_format starts with a dot
-        if not file_format.startswith('.'):
-            file_format = '.' + file_format
+        self.file_format = self.file_format_combo.currentText()
+        if not self.file_format.startswith('.'):
+            self.file_format = '.' + self.file_format
 
-        # Create output directory
         output_path = os.path.join(output_dir, folder_name)
-        try:
-            os.makedirs(output_path, exist_ok=True)
-        except Exception as e:
-            QMessageBox.critical(self,
-                                 "Error Creating Directory",
-                                 f"Failed to create output directory: {str(e)}")
-            return
-
-        # Get the list of images to process
+        os.makedirs(output_path, exist_ok=True)
+        
         images = self.get_selected_image_paths()
         if not images:
-            QMessageBox.warning(self,
-                                "No Images",
-                                "No images found in the project.")
+            QMessageBox.warning(self, "No Images", "No images found for processing.")
             return
-
-        # Collect annotation types to include
+            
         self.annotation_types = []
-        if self.patch_checkbox.isChecked():
+        if self.patch_checkbox.isChecked(): 
             self.annotation_types.append(PatchAnnotation)
-        if self.rectangle_checkbox.isChecked():
+        if self.rectangle_checkbox.isChecked(): 
             self.annotation_types.append(RectangleAnnotation)
         if self.polygon_checkbox.isChecked():
             self.annotation_types.append(PolygonAnnotation)
             self.annotation_types.append(MultiPolygonAnnotation)
 
-        # Create class mapping
-        self.class_mapping = {}
-
-        for i, label in enumerate(self.selected_labels):
-            # Leave 0 for background
-            self.class_mapping[label.short_label_code] = {
-                "label": label.to_dict(),
-                "index": i + 1
-            }
-
-        # Make the cursor busy
+        # --- Run Export Loop ---
         QApplication.setOverrideCursor(Qt.WaitCursor)
-        progress_bar = ProgressBar(self.annotation_window, "Exporting Segmentation Masks")
+        progress_bar = ProgressBar(self.annotation_window, "Exporting Masks")
         progress_bar.show()
         progress_bar.start_progress(len(images))
 
         try:
             for image_path in images:
-                # Create mask for this image
-                self.create_mask_for_image(image_path, output_path, file_format)
+                self.create_mask_for_image(image_path, output_path)
                 progress_bar.update_progress()
 
-            # Write the class mapping to a JSON file
-            self.export_class_mapping(output_path)
-
-            QMessageBox.information(self,
-                                    "Export Complete",
-                                    "Segmentation masks have been successfully exported")
+            self.export_metadata(output_path)
+            QMessageBox.information(self, "Export Complete", "Masks exported successfully.")
             self.accept()
-
         except Exception as e:
-            QMessageBox.critical(self,
-                                 "Error Exporting Masks",
-                                 f"An error occurred: {str(e)}")
-
+            QMessageBox.critical(self, "Error", f"An error occurred during export: {e}")
         finally:
-            # Make cursor normal again
             QApplication.restoreOverrideCursor()
-            progress_bar.finish_progress()
-            progress_bar.stop_progress()
             progress_bar.close()
 
-    def get_annotations_for_image(self, image_path):
-        """Get annotations for a specific image."""
-        # Get the selected labels' short label codes
-        selected_labels = [label.short_label_code for label in self.selected_labels]
-
-        # Get all annotations for this image
-        annotations = []
-
-        for annotation in self.annotation_window.get_image_annotations(image_path):
-            # Check that the annotation is of the correct type
-            if not isinstance(annotation, tuple(self.annotation_types)):
-                continue
-
-            # Check that the annotation's label is in the selected labels
-            if annotation.label.short_label_code not in selected_labels:
-                continue
-
-            # Add the annotation to the list based on its type, if selected
-            if self.patch_checkbox.isChecked() and isinstance(annotation, PatchAnnotation):
-                annotations.append(annotation)
-
-            elif self.rectangle_checkbox.isChecked() and isinstance(annotation, RectangleAnnotation):
-                annotations.append(annotation)
+    def create_mask_for_image(self, image_path, output_path):
+        height, width, has_georef, transform, crs = self.get_image_metadata(image_path, self.file_format)
+        if not height or not width:
+            print(f"Skipping {image_path}: could not determine dimensions.")
+            return
 
-            elif self.polygon_checkbox.isChecked() and isinstance(annotation, PolygonAnnotation):
-                annotations.append(annotation)
-                
-            elif self.polygon_checkbox.isChecked() and isinstance(annotation, MultiPolygonAnnotation):
-                for polygon in annotation.polygons:
-                    annotations.append(polygon)
+        # Initialize mask based on mode
+        if self.mask_mode == 'rgb':
+            mask = np.full((height, width, 3), self.background_value, dtype=np.uint8)
+        else:  # semantic or sfm
+            mask = np.full((height, width), self.background_value, dtype=np.uint8)
+
+        has_annotations = False
+        
+        # Draw annotations based on mode
+        if self.mask_mode in ['semantic', 'sfm']:
+            for label, value in self.labels_to_render:
+                annotations = self.get_annotations_for_image(image_path, label)
+                if annotations:
+                    has_annotations = True
+                    self.draw_annotations_on_mask(mask, annotations, value)
+        elif self.mask_mode == 'rgb':
+            for label, color in self.labels_to_render:
+                annotations = self.get_annotations_for_image(image_path, label)
+                if annotations:
+                    has_annotations = True
+                    self.draw_annotations_on_mask(mask, annotations, color)
+        
+        if not has_annotations and not self.include_negative_samples_checkbox.isChecked():
+            return
 
+        # Always save as PNG
+        filename = f"{os.path.splitext(os.path.basename(image_path))[0]}.png"
+        mask_path = os.path.join(output_path, filename)
+
+        if self.mask_mode == 'rgb':
+            # OpenCV expects BGR, so convert from RGB
+            mask = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(mask_path, mask)
+
+    def export_metadata(self, output_path):
+        if self.mask_mode == 'semantic':
+            class_mapping = {}
+            if self.label_table.cellWidget(0, 0).findChild(QCheckBox).isChecked():
+                background_label = "Background"
+                background_index = self.label_table.cellWidget(0, 2).value()
+                class_mapping[background_label] = {
+                    "label": background_label,
+                    "index": background_index
+                }
+            
+            for label, value in self.labels_to_render:
+                class_mapping[label.short_label_code] = {"label": label.to_dict(), "index": value}
+            
+            with open(os.path.join(output_path, "class_mapping.json"), 'w') as f:
+                json.dump(class_mapping, f, indent=4)
+        
+        elif self.mask_mode == 'rgb':
+            color_legend = {}
+            if self.label_table.cellWidget(0, 0).findChild(QCheckBox).isChecked():
+                color_legend["Background"] = self.background_value
+
+            for label, color in self.labels_to_render:
+                color_legend[label.short_label_code] = color
+            
+            with open(os.path.join(output_path, "color_legend.json"), 'w') as f:
+                json.dump(color_legend, f, indent=4)
+        
+        # No metadata file needed for SfM mode
+
+    def get_annotations_for_image(self, image_path, label):
+        annotations = []
+        for ann in self.annotation_window.get_image_annotations(image_path):
+            if ann.label.short_label_code == label.short_label_code and isinstance(ann, tuple(self.annotation_types)):
+                if isinstance(ann, MultiPolygonAnnotation):
+                    annotations.extend(ann.polygons)
+                else:
+                    annotations.append(ann)
         return annotations
 
+    def draw_annotations_on_mask(self, mask, annotations, value):
+        for ann in annotations:
+            if isinstance(ann, (PatchAnnotation, RectangleAnnotation)):
+                p1 = (int(ann.top_left.x()), int(ann.top_left.y()))
+                p2 = (int(ann.bottom_right.x()), int(ann.bottom_right.y()))
+                cv2.rectangle(mask, p1, p2, value, -1)
+            elif isinstance(ann, PolygonAnnotation):
+                points = np.array([[p.x(), p.y()] for p in ann.points], dtype=np.int32)
+                cv2.fillPoly(mask, [points], value)
+
     def get_image_metadata(self, image_path, file_format):
-        """Get the dimensions of the image, and check for georeferencing."""
-        # Check if image has georeferencing that needs to be preserved
-        transform = None
-        crs = None
-        has_georef = False
-        height = None
-        width = None
-
-        # Get the raster from the raster manager
+        transform, crs, has_georef = None, None, False
+        width, height = None, None
         raster = self.image_window.raster_manager.get_raster(image_path)
-
-        # Only check for georeferencing if using TIF format and checkbox is checked
-        can_preserve_georef = self.preserve_georef_checkbox.isChecked() and file_format.lower() == '.tif'
+        can_preserve = self.preserve_georef_checkbox.isChecked() and file_format.lower() == '.tif'
 
         if raster and raster.rasterio_src:
-            # Get dimensions from the raster
-            width = raster.width
-            height = raster.height
-
-            # Check for georeferencing if needed
-            if can_preserve_georef and hasattr(raster.rasterio_src, 'transform'):
+            width, height = raster.width, raster.height
+            if can_preserve and hasattr(raster.rasterio_src, 'transform'):
                 transform = raster.rasterio_src.transform
                 if transform and not transform.is_identity:
                     crs = raster.rasterio_src.crs
                     has_georef = True
         else:
-            # Fallback to direct file access if raster is not available
             try:
-                if can_preserve_georef:
+                if can_preserve:
                     with rasterio.open(image_path) as src:
-                        if src.transform and not src.transform.is_identity:
-                            transform = src.transform
-                            crs = src.crs
-                            has_georef = True
                         width, height = src.width, src.height
+                        if src.transform and not src.transform.is_identity:
+                            transform, crs, has_georef = src.transform, src.crs, True
                 else:
-                    # Use PIL for non-georeferenced images
-                    image = Image.open(image_path)
-                    width, height = image.size
+                    with Image.open(image_path) as img:
+                        width, height = img.size
             except Exception as e:
-                print(f"Error loading image {image_path}: {str(e)}")
-
+                print(f"Error reading metadata for {image_path}: {e}")
         return height, width, has_georef, transform, crs
 
-    def draw_annotations_on_mask(self, mask, annotations):
-        """Draw annotations on the mask."""
-        # Draw each annotation on the mask
-        for annotation in annotations:
-            # Get the label index for the annotation
-            short_label_code = annotation.label.short_label_code
-            label_index = self.class_mapping[short_label_code]["index"]
-
-            # Draw the patch annotation
-            if isinstance(annotation, PatchAnnotation):
-                # Draw a filled rectangle
-                cv2.rectangle(mask,
-                              (int(annotation.center_xy.x() - annotation.annotation_size / 2),
-                               int(annotation.center_xy.y() - annotation.annotation_size / 2)),
-                              (int(annotation.center_xy.x() + annotation.annotation_size / 2),
-                               int(annotation.center_xy.y() + annotation.annotation_size / 2)),
-                              label_index, -1)  # -1 means filled
-
-            # Draw the rectangle annotation
-            elif isinstance(annotation, RectangleAnnotation):
-                # Draw a filled rectangle
-                cv2.rectangle(mask,
-                              (int(annotation.top_left.x()), int(annotation.top_left.y())),
-                              (int(annotation.bottom_right.x()), int(annotation.bottom_right.y())),
-                              label_index, -1)  # -1 means filled
-
-            # Draw the polygon annotation
-            elif isinstance(annotation, PolygonAnnotation):
-                # Draw a filled polygon
-                points = np.array([[p.x(), p.y()] for p in annotation.points]).astype(np.int32)
-                cv2.fillPoly(mask, [points], label_index)
-                
-            # Draw the multipolygon annotation
-            elif isinstance(annotation, MultiPolygonAnnotation):
-                for polygon in annotation.polygons:
-                    points = np.array([[p.x(), p.y()] for p in polygon.points]).astype(np.int32)
-                    cv2.fillPoly(mask, [points], label_index)
-
-        return mask
-
-    def create_mask_for_image(self, image_path, output_path, file_format):
-        """Create a segmentation mask for a single image"""
-        # Get the annotations for this image
-        annotations = self.get_annotations_for_image(image_path)
-
-        if not annotations and not self.include_negative_samples_checkbox.isChecked():
-            return  # Skip images with no annotations if the user doesn't want negative samples
-
-        # Get the image dimensions, georeferencing
-        height, width, has_georef, transform, crs = self.get_image_metadata(image_path, file_format)
+    # --- Row Movement and UI Helpers ---
+    def move_row_up(self):
+        current_row = self.label_table.currentRow()
+        if current_row > 0:
+            self.swap_rows(current_row, current_row - 1)
+            self.label_table.selectRow(current_row - 1)
+
+    def move_row_down(self):
+        current_row = self.label_table.currentRow()
+        if 0 <= current_row < self.label_table.rowCount() - 1:
+            self.swap_rows(current_row, current_row + 1)
+            self.label_table.selectRow(current_row + 1)
+
+    def swap_rows(self, row1, row2):
+        # Because widgets are complex to swap, we just repopulate the table
+        # while preserving the core data (label order). This is simpler and more robust.
+        
+        # Step 1: Extract the core data and checkbox state from the table
+        table_data = []
+        for r in range(self.label_table.rowCount()):
+            is_checked = self.label_table.cellWidget(r, 0).findChild(QCheckBox).isChecked()
+            label_code = self.label_table.item(r, 1).data(Qt.UserRole)
+            table_data.append({'code': label_code, 'checked': is_checked})
+
+        # Step 2: Swap the data for the two rows
+        table_data[row1], table_data[row2] = table_data[row2], table_data[row1]
+
+        # Step 3: Map old label order to new order
+        new_label_order = [data['code'] for data in table_data if data['code'] != 'Background']
+        
+        def label_sort_key(x):
+            if x.short_label_code in new_label_order:
+                return new_label_order.index(x.short_label_code)
+            else:
+                return float('inf')
+        self.label_window.labels.sort(key=label_sort_key)
 
-        if not height or not width:
-            print(f"Could not get dimensions for image: {image_path}")
-            return
+        # Step 4: Repopulate and restore checkbox states
+        self.populate_label_table()
+        for r, data in enumerate(table_data):
+            self.label_table.cellWidget(r, 0).findChild(QCheckBox).setChecked(data['checked'])
 
-        # Create a blank mask
-        mask = np.zeros((height, width), dtype=np.uint8)
-
-        # Draw annotations on the mask
-        mask = self.draw_annotations_on_mask(mask, annotations)
-
-        # Save the mask
-        filename = os.path.basename(image_path)
-        name_without_ext = os.path.splitext(filename)[0]
-        mask_filename = f"{name_without_ext}{file_format}"
-        mask_path = os.path.join(output_path, mask_filename)
-
-        # If we have georeferencing and we need to preserve it, use rasterio to save
-        if has_georef and file_format.lower() == '.tif':
-            with rasterio.open(
-                mask_path,
-                'w',
-                driver='GTiff',
-                height=height,
-                width=width,
-                count=1,
-                dtype=mask.dtype,
-                crs=crs,
-                transform=transform,
-                compress='lzw',
-            ) as dst:
-                dst.write(mask, 1)
+    def update_georef_availability(self):
+        is_tif = '.tif' in self.file_format_combo.currentText().lower()
+        self.preserve_georef_checkbox.setEnabled(is_tif)
+        if not is_tif:
+            self.preserve_georef_checkbox.setChecked(False)
+            self.georef_note.setStyleSheet("color: red; font-style: italic;")
         else:
-            # Use OpenCV as before for non-georeferenced images
-            cv2.imwrite(mask_path, mask)
-
-        if not os.path.exists(mask_path):
-            print(f"Warning: Failed to save mask to {mask_path}")
+            self.georef_note.setStyleSheet("color: #666; font-style: italic;")
 
     def closeEvent(self, event):
-        """Handle the close event."""
-        # Clean up any resources if needed
         super().closeEvent(event)
+