lazylabel-gui 1.0.7__py3-none-any.whl → 1.0.9__py3-none-any.whl

lazylabel/main.py

@@ -3,6 +3,7 @@ import os
 import numpy as np
 import qdarktheme
 import cv2
+import json
 from PyQt6.QtWidgets import (
     QApplication,
     QMainWindow,
@@ -15,6 +16,8 @@ from PyQt6.QtWidgets import (
     QTableWidgetItem,
     QGraphicsPolygonItem,
     QTableWidgetSelectionRange,
+    QSpacerItem,
+    QHeaderView,
 )
 from PyQt6.QtGui import (
     QPixmap,
@@ -60,14 +63,15 @@ class MainWindow(QMainWindow):
         self.current_file_index = QModelIndex()
 
         self.next_class_id = 0
+        self.class_aliases = {}
 
-        self.class_aliases = {}  # {class_id: "alias_string"}
+        self._original_point_radius = 0.3
+        self._original_line_thickness = 0.5
+        self.point_radius = self._original_point_radius
+        self.line_thickness = self._original_line_thickness
 
-        self.point_radius = 0.3
-        self.line_thickness = 0.5
-
-        self._original_point_radius = self.point_radius
-        self._original_line_thickness = self.line_thickness
+        self.pan_multiplier = 1.0
+        self.polygon_join_threshold = 2
 
         self.point_items, self.positive_points, self.negative_points = [], [], []
         self.polygon_points, self.polygon_preview_items = [], []
@@ -87,6 +91,10 @@ class MainWindow(QMainWindow):
         self.file_model = CustomFileSystemModel()
         self.right_panel.file_tree.setModel(self.file_model)
         self.right_panel.file_tree.setColumnWidth(0, 200)
+        file_tree = self.right_panel.file_tree
+        header = file_tree.header()
+        header.setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)
+        header.setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
 
         main_layout = QHBoxLayout()
         main_layout.addWidget(self.control_panel)
@@ -101,6 +109,7 @@ class MainWindow(QMainWindow):
         )
         self.setup_connections()
         self.set_sam_mode()
+        self.set_annotation_size(10)
 
     def setup_connections(self):
         self._original_mouse_press = self.viewer.scene().mousePressEvent
@@ -122,7 +131,6 @@ class MainWindow(QMainWindow):
         self.right_panel.segment_table.itemSelectionChanged.connect(
             self.highlight_selected_segments
         )
-        self.right_panel.segment_table.itemChanged.connect(self.handle_class_id_change)
         self.right_panel.class_table.itemChanged.connect(self.handle_alias_change)
         self.right_panel.btn_reassign_classes.clicked.connect(self.reassign_class_ids)
         self.right_panel.class_filter_combo.currentIndexChanged.connect(
@@ -137,11 +145,131 @@ class MainWindow(QMainWindow):
         self.control_panel.btn_clear_points.clicked.connect(self.clear_all_points)
         self.control_panel.btn_fit_view.clicked.connect(self.viewer.fitInView)
 
-        # **FIX:** Use QShortcut for reliable global hotkeys
-        next_shortcut = QShortcut(QKeySequence(Qt.Key.Key_Right), self)
-        next_shortcut.activated.connect(self.load_next_image)
-        prev_shortcut = QShortcut(QKeySequence(Qt.Key.Key_Left), self)
-        prev_shortcut.activated.connect(self.load_previous_image)
+        self.control_panel.size_slider.valueChanged.connect(self.set_annotation_size)
+        self.control_panel.pan_slider.valueChanged.connect(self.set_pan_multiplier)
+        self.control_panel.join_slider.valueChanged.connect(
+            self.set_polygon_join_threshold
+        )
+
+        self.control_panel.chk_save_npz.stateChanged.connect(
+            self.handle_save_checkbox_change
+        )
+        self.control_panel.chk_save_txt.stateChanged.connect(
+            self.handle_save_checkbox_change
+        )
+
+        self.control_panel.btn_toggle_visibility.clicked.connect(self.toggle_left_panel)
+        self.right_panel.btn_toggle_visibility.clicked.connect(self.toggle_right_panel)
+
+        QShortcut(QKeySequence(Qt.Key.Key_Right), self, self.load_next_image)
+        QShortcut(QKeySequence(Qt.Key.Key_Left), self, self.load_previous_image)
+        QShortcut(QKeySequence(Qt.Key.Key_1), self, self.set_sam_mode)
+        QShortcut(QKeySequence(Qt.Key.Key_2), self, self.set_polygon_mode)
+        QShortcut(QKeySequence(Qt.Key.Key_E), self, self.toggle_selection_mode)
+        QShortcut(QKeySequence(Qt.Key.Key_Q), self, self.toggle_pan_mode)
+        QShortcut(QKeySequence(Qt.Key.Key_R), self, self.toggle_edit_mode)
+        QShortcut(QKeySequence(Qt.Key.Key_C), self, self.clear_all_points)
+        QShortcut(QKeySequence(Qt.Key.Key_Escape), self, self.handle_escape_press)
+        QShortcut(QKeySequence(Qt.Key.Key_V), self, self.delete_selected_segments)
+        QShortcut(
+            QKeySequence(Qt.Key.Key_Backspace), self, self.delete_selected_segments
+        )
+        QShortcut(QKeySequence(Qt.Key.Key_M), self, self.handle_merge_press)
+        QShortcut(QKeySequence("Ctrl+Z"), self, self.undo_last_action)
+        QShortcut(
+            QKeySequence("Ctrl+A"), self, self.right_panel.segment_table.selectAll
+        )
+        QShortcut(QKeySequence(Qt.Key.Key_Space), self, self.handle_space_press)
+        QShortcut(QKeySequence(Qt.Key.Key_Return), self, self.handle_enter_press)
+        QShortcut(QKeySequence(Qt.Key.Key_Enter), self, self.handle_enter_press)
+        QShortcut(QKeySequence(Qt.Key.Key_Period), self, self.viewer.fitInView)
+
+    def toggle_left_panel(self):
+        is_visible = self.control_panel.main_controls_widget.isVisible()
+        self.control_panel.main_controls_widget.setVisible(not is_visible)
+        if is_visible:
+            self.control_panel.btn_toggle_visibility.setText("> Show")
+            self.control_panel.setFixedWidth(
+                self.control_panel.btn_toggle_visibility.sizeHint().width() + 20
+            )
+        else:
+            self.control_panel.btn_toggle_visibility.setText("< Hide")
+            self.control_panel.setFixedWidth(250)
+
+    def toggle_right_panel(self):
+        is_visible = self.right_panel.main_controls_widget.isVisible()
+        self.right_panel.main_controls_widget.setVisible(not is_visible)
+        layout = self.right_panel.v_layout
+
+        if is_visible:  # Content is now hidden
+            layout.addStretch(1)
+            self.right_panel.btn_toggle_visibility.setText("< Show")
+            self.right_panel.setFixedWidth(
+                self.right_panel.btn_toggle_visibility.sizeHint().width() + 20
+            )
+        else:  # Content is now visible
+            # Remove the stretch so the content can expand
+            for i in range(layout.count()):
+                item = layout.itemAt(i)
+                if isinstance(item, QSpacerItem):
+                    layout.removeItem(item)
+                    break
+            self.right_panel.btn_toggle_visibility.setText("Hide >")
+            self.right_panel.setFixedWidth(350)
+
+    def handle_save_checkbox_change(self):
+        is_npz_checked = self.control_panel.chk_save_npz.isChecked()
+        is_txt_checked = self.control_panel.chk_save_txt.isChecked()
+
+        if not is_npz_checked and not is_txt_checked:
+            sender = self.sender()
+            if sender == self.control_panel.chk_save_npz:
+                self.control_panel.chk_save_txt.setChecked(True)
+            else:
+                self.control_panel.chk_save_npz.setChecked(True)
+
+    def set_annotation_size(self, value):
+        multiplier = value / 10.0
+        self.point_radius = self._original_point_radius * multiplier
+        self.line_thickness = self._original_line_thickness * multiplier
+
+        self.control_panel.size_label.setText(f"Annotation Size: {multiplier:.1f}x")
+
+        if self.control_panel.size_slider.value() != value:
+            self.control_panel.size_slider.setValue(value)
+
+        self.display_all_segments()
+        self.clear_all_points()
+
+    def set_pan_multiplier(self, value):
+        self.pan_multiplier = value / 10.0
+        self.control_panel.pan_label.setText(f"Pan Speed: {self.pan_multiplier:.1f}x")
+
+    def set_polygon_join_threshold(self, value):
+        self.polygon_join_threshold = value
+        self.control_panel.join_label.setText(f"Polygon Join Distance: {value}px")
+
+    def handle_escape_press(self):
+        self.right_panel.segment_table.clearSelection()
+        self.right_panel.class_table.clearSelection()
+        self.clear_all_points()
+        self.viewer.setFocus()
+
+    def handle_space_press(self):
+        if self.mode == "polygon" and self.polygon_points:
+            self.finalize_polygon()
+        else:
+            self.save_current_segment()
+
+    def handle_enter_press(self):
+        if self.mode == "polygon" and self.polygon_points:
+            self.finalize_polygon()
+        else:
+            self.save_output_to_npz()
+
+    def handle_merge_press(self):
+        self.assign_selected_to_class()
+        self.right_panel.segment_table.clearSelection()
 
     def show_notification(self, message, duration=3000):
         self.control_panel.notification_label.setText(message)
@@ -152,10 +280,8 @@ class MainWindow(QMainWindow):
     def _get_color_for_class(self, class_id):
         if class_id is None:
             return QColor.fromHsv(0, 0, 128)
-
         hue = int((class_id * 222.4922359) % 360)
         color = QColor.fromHsv(hue, 220, 220)
-
         if not color.isValid():
             return QColor(Qt.GlobalColor.white)
         return color
@@ -256,6 +382,7 @@ class MainWindow(QMainWindow):
                 self.reset_state()
                 self.viewer.set_photo(pixmap)
                 self.sam_model.set_image(self.current_image_path)
+                self.load_class_aliases()
                 self.load_existing_mask()
                 self.right_panel.file_tree.setCurrentIndex(index)
                 self.viewer.setFocus()
@@ -288,9 +415,8 @@ class MainWindow(QMainWindow):
 
     def reset_state(self):
         self.clear_all_points()
-        # Preserve aliases between images in the same session
-        # self.class_aliases.clear()
         self.segments.clear()
+        self.class_aliases.clear()
         self.next_class_id = 0
         self.update_all_lists()
         items_to_remove = [
@@ -305,83 +431,55 @@ class MainWindow(QMainWindow):
 
     def keyPressEvent(self, event):
         key, mods = event.key(), event.modifiers()
-        if event.isAutoRepeat():
+
+        if event.isAutoRepeat() and key not in {
+            Qt.Key.Key_W,
+            Qt.Key.Key_A,
+            Qt.Key.Key_S,
+            Qt.Key.Key_D,
+        }:
             return
 
-        pan_multiplier = 5.0 if (mods & Qt.KeyboardModifier.ShiftModifier) else 2
+        shift_multiplier = 5.0 if mods & Qt.KeyboardModifier.ShiftModifier else 1.0
 
         if key == Qt.Key.Key_W:
-            amount = int(self.viewer.height() * 0.1 * pan_multiplier)
+            amount = int(
+                self.viewer.height() * 0.1 * self.pan_multiplier * shift_multiplier
+            )
             self.viewer.verticalScrollBar().setValue(
                 self.viewer.verticalScrollBar().value() - amount
             )
         elif key == Qt.Key.Key_S:
-            amount = int(self.viewer.height() * 0.1 * pan_multiplier)
+            amount = int(
+                self.viewer.height() * 0.1 * self.pan_multiplier * shift_multiplier
+            )
             self.viewer.verticalScrollBar().setValue(
                 self.viewer.verticalScrollBar().value() + amount
             )
-        elif key == Qt.Key.Key_A and not (mods & Qt.KeyboardModifier.ControlModifier):
-            amount = int(self.viewer.width() * 0.1 * pan_multiplier)
+        elif key == Qt.Key.Key_A:
+            amount = int(
+                self.viewer.width() * 0.1 * self.pan_multiplier * shift_multiplier
+            )
             self.viewer.horizontalScrollBar().setValue(
                 self.viewer.horizontalScrollBar().value() - amount
             )
         elif key == Qt.Key.Key_D:
-            amount = int(self.viewer.width() * 0.1 * pan_multiplier)
+            amount = int(
+                self.viewer.width() * 0.1 * self.pan_multiplier * shift_multiplier
+            )
             self.viewer.horizontalScrollBar().setValue(
                 self.viewer.horizontalScrollBar().value() + amount
             )
-        elif key == Qt.Key.Key_Period:
-            self.viewer.fitInView()
-        # Other keybindings
-        elif key == Qt.Key.Key_1:
-            self.set_sam_mode()
-        elif key == Qt.Key.Key_2:
-            self.set_polygon_mode()
-        elif key == Qt.Key.Key_E:
-            self.toggle_selection_mode()
-        elif key == Qt.Key.Key_Q:
-            self.toggle_pan_mode()
-        elif key == Qt.Key.Key_R:
-            self.toggle_edit_mode()
-        elif key == Qt.Key.Key_C or key == Qt.Key.Key_Escape:
-            self.clear_all_points()
-        elif key == Qt.Key.Key_V or key == Qt.Key.Key_Backspace:
-            self.delete_selected_segments()
-        elif key == Qt.Key.Key_M:
-            self.assign_selected_to_class()
-            self.right_panel.segment_table.clearSelection()
-        elif key == Qt.Key.Key_Z and mods == Qt.KeyboardModifier.ControlModifier:
-            self.undo_last_action()
-        elif key == Qt.Key.Key_A and mods == Qt.KeyboardModifier.ControlModifier:
-            self.right_panel.segment_table.selectAll()
-        elif key == Qt.Key.Key_Space:
-            if self.mode == "polygon" and self.polygon_points:
-                self.finalize_polygon()
-            else:
-                self.save_current_segment()
-        elif key == Qt.Key.Key_Return or key == Qt.Key.Key_Enter:
-            if self.mode == "polygon" and self.polygon_points:
-                self.finalize_polygon()
-            else:
-                self.save_output_to_npz()
         elif (
             key == Qt.Key.Key_Equal or key == Qt.Key.Key_Plus
         ) and mods == Qt.KeyboardModifier.ControlModifier:
-            self.point_radius = min(20, self.point_radius + self._original_point_radius)
-            self.line_thickness = min(
-                20, self.line_thickness + self._original_line_thickness
-            )
-            self.display_all_segments()
-            self.clear_all_points()
+            current_val = self.control_panel.size_slider.value()
+            self.control_panel.size_slider.setValue(current_val + 1)
         elif key == Qt.Key.Key_Minus and mods == Qt.KeyboardModifier.ControlModifier:
-            self.point_radius = max(
-                0.3, self.point_radius - self._original_point_radius
-            )
-            self.line_thickness = max(
-                0.5, self.line_thickness - self._original_line_thickness
-            )
-            self.display_all_segments()
-            self.clear_all_points()
+            current_val = self.control_panel.size_slider.value()
+            self.control_panel.size_slider.setValue(current_val - 1)
+        else:
+            super().keyPressEvent(event)
 
     def scene_mouse_press(self, event):
         self._original_mouse_press(event)
@@ -432,15 +530,12 @@ class MainWindow(QMainWindow):
         elif self.mode == "polygon" and self.polygon_points:
             if self.rubber_band_line is None:
                 self.rubber_band_line = QGraphicsLineItem()
-
                 line_color = QColor(Qt.GlobalColor.white)
                 line_color.setAlpha(150)
-
                 self.rubber_band_line.setPen(
                     QPen(line_color, self.line_thickness, Qt.PenStyle.DotLine)
                 )
                 self.viewer.scene().addItem(self.rubber_band_line)
-
             self.rubber_band_line.setLine(
                 self.polygon_points[-1].x(),
                 self.polygon_points[-1].y(),
@@ -454,7 +549,6 @@ class MainWindow(QMainWindow):
     def scene_mouse_release(self, event):
         if self.mode == "pan":
             self.viewer.set_cursor(Qt.CursorShape.OpenHandCursor)
-
         if self.mode == "edit" and self.is_dragging_polygon:
             self.is_dragging_polygon = False
             self.drag_initial_vertices.clear()
@@ -463,12 +557,10 @@ class MainWindow(QMainWindow):
     def undo_last_action(self):
         if self.mode == "polygon" and self.polygon_points:
             self.polygon_points.pop()
-
             for item in self.polygon_preview_items:
                 if item.scene():
                     self.viewer.scene().removeItem(item)
             self.polygon_preview_items.clear()
-
             for point in self.polygon_points:
                 point_diameter = self.point_radius * 2
                 point_color = QColor(Qt.GlobalColor.blue)
@@ -483,21 +575,17 @@ class MainWindow(QMainWindow):
                 dot.setPen(QPen(Qt.GlobalColor.transparent))
                 self.viewer.scene().addItem(dot)
                 self.polygon_preview_items.append(dot)
-
             self.draw_polygon_preview()
-
         elif self.mode == "sam_points" and self.point_items:
             item_to_remove = self.point_items.pop()
             point_pos = item_to_remove.rect().topLeft() + QPointF(
                 self.point_radius, self.point_radius
             )
             point_coords = [int(point_pos.x()), int(point_pos.y())]
-
             if point_coords in self.positive_points:
                 self.positive_points.remove(point_coords)
             elif point_coords in self.negative_points:
                 self.negative_points.remove(point_coords)
-
             self.viewer.scene().removeItem(item_to_remove)
             self.update_segmentation()
 
@@ -574,14 +662,13 @@ class MainWindow(QMainWindow):
         if existing_class_ids:
             target_class_id = min(existing_class_ids)
         else:
-            target_class_id = self.segments[selected_indices[0]].get("class_id")
+            target_class_id = self.next_class_id
 
         for i in selected_indices:
             self.segments[i]["class_id"] = target_class_id
 
         self._update_next_class_id()
         self.update_all_lists()
-        self.right_panel.segment_table.clearSelection()
         self.viewer.setFocus()
 
     def rasterize_polygon(self, vertices):
@@ -606,7 +693,6 @@ class MainWindow(QMainWindow):
         for i, seg_dict in enumerate(self.segments):
             self.segment_items[i] = []
             class_id = seg_dict.get("class_id")
-
             base_color = self._get_color_for_class(class_id)
 
             if seg_dict["type"] == "Polygon":
@@ -621,7 +707,6 @@ class MainWindow(QMainWindow):
                 poly_item.setPen(QPen(Qt.GlobalColor.transparent))
                 self.viewer.scene().addItem(poly_item)
                 self.segment_items[i].append(poly_item)
-
                 base_color.setAlpha(150)
                 vertex_color = QBrush(base_color)
                 point_diameter = self.point_radius * 2
@@ -658,7 +743,6 @@ class MainWindow(QMainWindow):
                 hover_pixmap = mask_to_pixmap(
                     seg_dict["mask"], base_color.getRgb()[:3], alpha=170
                 )
-
                 pixmap_item = HoverablePixmapItem()
                 pixmap_item.set_pixmaps(default_pixmap, hover_pixmap)
                 self.viewer.scene().addItem(pixmap_item)
@@ -697,7 +781,7 @@ class MainWindow(QMainWindow):
                 self.highlight_items.append(highlight_item)
 
     def update_all_lists(self):
-        self.update_class_list()  # Must be before filter combo
+        self.update_class_list()
         self.update_class_filter_combo()
         self.update_segment_table()
         self.display_all_segments()
@@ -727,19 +811,24 @@ class MainWindow(QMainWindow):
         for row, (original_index, seg) in enumerate(display_segments):
             class_id = seg.get("class_id")
             color = self._get_color_for_class(class_id)
-
             class_id_str = str(class_id) if class_id is not None else "N/A"
+
+            alias_str = "N/A"
+            if class_id is not None:
+                alias_str = self.class_aliases.get(class_id, str(class_id))
+            alias_item = QTableWidgetItem(alias_str)
+
             index_item = NumericTableWidgetItem(str(original_index + 1))
             class_item = NumericTableWidgetItem(class_id_str)
-            type_item = QTableWidgetItem(seg.get("type", "N/A"))
 
             index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
-            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
+            class_item.setFlags(class_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
+            alias_item.setFlags(alias_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
             index_item.setData(Qt.ItemDataRole.UserRole, original_index)
 
             table.setItem(row, 0, index_item)
             table.setItem(row, 1, class_item)
-            table.setItem(row, 2, type_item)
+            table.setItem(row, 2, alias_item)
 
             for col in range(table.columnCount()):
                 if table.item(row, col):
@@ -759,19 +848,7 @@ class MainWindow(QMainWindow):
         class_table = self.right_panel.class_table
         class_table.blockSignals(True)
 
-        # Preserve existing aliases during update
-        current_aliases = {}
-        for row in range(class_table.rowCount()):
-            try:
-                alias = class_table.item(row, 0).text()
-                cid = int(class_table.item(row, 1).text())
-                current_aliases[cid] = alias
-            except (AttributeError, ValueError):
-                continue
-        self.class_aliases.update(current_aliases)
-
-        class_table.clearContents()
-
+        preserved_aliases = self.class_aliases.copy()
         unique_class_ids = sorted(
             list(
                 {
@@ -781,19 +858,22 @@ class MainWindow(QMainWindow):
                 }
             )
         )
-        class_table.setRowCount(len(unique_class_ids))
 
+        new_aliases = {}
+        for cid in unique_class_ids:
+            new_aliases[cid] = preserved_aliases.get(cid, str(cid))
+
+        self.class_aliases = new_aliases
+
+        class_table.clearContents()
+        class_table.setRowCount(len(unique_class_ids))
         for row, cid in enumerate(unique_class_ids):
-            alias = self.class_aliases.get(cid, str(cid))
-            alias_item = QTableWidgetItem(alias)
+            alias_item = QTableWidgetItem(self.class_aliases.get(cid))
             id_item = QTableWidgetItem(str(cid))
-
             id_item.setFlags(id_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
-
             color = self._get_color_for_class(cid)
             alias_item.setBackground(QBrush(color))
             id_item.setBackground(QBrush(color))
-
             class_table.setItem(row, 0, alias_item)
             class_table.setItem(row, 1, id_item)
 
@@ -810,7 +890,6 @@ class MainWindow(QMainWindow):
                 }
             )
         )
-
         current_selection = combo.currentText()
         combo.blockSignals(True)
         combo.clear()
@@ -821,7 +900,6 @@ class MainWindow(QMainWindow):
                 for cid in unique_class_ids
             ]
         )
-
         if combo.findText(current_selection) > -1:
             combo.setCurrentText(current_selection)
         else:
@@ -830,7 +908,6 @@ class MainWindow(QMainWindow):
 
     def reassign_class_ids(self):
         class_table = self.right_panel.class_table
-
         ordered_ids = []
         for row in range(class_table.rowCount()):
             id_item = class_table.item(row, 1)
@@ -839,75 +916,37 @@ class MainWindow(QMainWindow):
                     ordered_ids.append(int(id_item.text()))
                 except ValueError:
                     continue
-
         id_map = {old_id: new_id for new_id, old_id in enumerate(ordered_ids)}
-
         for seg in self.segments:
             old_id = seg.get("class_id")
             if old_id in id_map:
                 seg["class_id"] = id_map[old_id]
-
         new_aliases = {
             id_map[old_id]: self.class_aliases.get(old_id, str(old_id))
             for old_id in ordered_ids
             if old_id in self.class_aliases
         }
         self.class_aliases = new_aliases
-
         self._update_next_class_id()
         self.update_all_lists()
         self.viewer.setFocus()
 
     def handle_alias_change(self, item):
-        if item.column() != 0:  # Alias column
+        if item.column() != 0:
             return
-
         class_table = self.right_panel.class_table
         class_table.blockSignals(True)
-
         id_item = class_table.item(item.row(), 1)
         if id_item:
             try:
                 class_id = int(id_item.text())
                 self.class_aliases[class_id] = item.text()
             except (ValueError, AttributeError):
-                pass  # Ignore if ID item is not valid
-
+                pass
         class_table.blockSignals(False)
-        self.update_class_filter_combo()  # Refresh filter to show new alias
-
-    def handle_class_id_change(self, item):
-        if item.column() != 1:  # Class ID column in segment table
-            return
-        table = self.right_panel.segment_table
-        index_item = table.item(item.row(), 0)
-        if not index_item:
-            return
-
-        table.blockSignals(True)
-        try:
-            new_class_id_text = item.text()
-            if not new_class_id_text.strip():
-                raise ValueError("Class ID cannot be empty.")
-            new_class_id = int(new_class_id_text)
-            original_index = index_item.data(Qt.ItemDataRole.UserRole)
 
-            if original_index is None or original_index >= len(self.segments):
-                raise IndexError("Invalid segment index found in table.")
-
-            self.segments[original_index]["class_id"] = new_class_id
-            self._update_next_class_id()
-            self.update_all_lists()
-        except (ValueError, TypeError, AttributeError, IndexError) as e:
-            original_index = index_item.data(Qt.ItemDataRole.UserRole)
-            if original_index is not None and original_index < len(self.segments):
-                original_class_id = self.segments[original_index].get("class_id")
-                item.setText(
-                    str(original_class_id) if original_class_id is not None else "N/A"
-                )
-        finally:
-            table.blockSignals(False)
-            self.viewer.setFocus()
+        self.update_class_filter_combo()
+        self.update_segment_table()
 
     def get_selected_segment_indices(self):
         table = self.right_panel.segment_table
@@ -920,95 +959,136 @@ class MainWindow(QMainWindow):
         ]
 
     def save_output_to_npz(self):
-        if not self.segments or not self.current_image_path:
+        save_npz = self.control_panel.chk_save_npz.isChecked()
+        save_txt = self.control_panel.chk_save_txt.isChecked()
+        save_aliases = self.control_panel.chk_save_class_aliases.isChecked()
+
+        if not self.current_image_path or not any([save_npz, save_txt, save_aliases]):
             return
+
         self.right_panel.status_label.setText("Saving...")
         QApplication.processEvents()
 
-        output_path = os.path.splitext(self.current_image_path)[0] + ".npz"
-        h, w = (
-            self.viewer._pixmap_item.pixmap().height(),
-            self.viewer._pixmap_item.pixmap().width(),
-        )
-
-        class_table = self.right_panel.class_table
-        ordered_ids = [
-            int(class_table.item(row, 1).text())
-            for row in range(class_table.rowCount())
-            if class_table.item(row, 1) is not None
-        ]
-
-        if not ordered_ids:
-            self.right_panel.status_label.setText("Save failed: No classes defined.")
-            QTimer.singleShot(3000, lambda: self.right_panel.status_label.clear())
-            return
-
-        id_map = {old_id: new_id for new_id, old_id in enumerate(ordered_ids)}
-        num_final_classes = len(ordered_ids)
-        final_mask_tensor = np.zeros((h, w, num_final_classes), dtype=np.uint8)
+        saved_something = False
 
-        for seg in self.segments:
-            class_id = seg.get("class_id")
-            if class_id not in id_map:
-                continue
-            new_channel_idx = id_map[class_id]
-            mask = (
-                self.rasterize_polygon(seg["vertices"])
-                if seg["type"] == "Polygon"
-                else seg.get("mask")
-            )
-            if mask is not None:
-                final_mask_tensor[:, :, new_channel_idx] = np.logical_or(
-                    final_mask_tensor[:, :, new_channel_idx], mask
+        if save_npz or save_txt:
+            if not self.segments:
+                self.show_notification("No segments to save.")
+            else:
+                h, w = (
+                    self.viewer._pixmap_item.pixmap().height(),
+                    self.viewer._pixmap_item.pixmap().width(),
                 )
+                class_table = self.right_panel.class_table
+                ordered_ids = [
+                    int(class_table.item(row, 1).text())
+                    for row in range(class_table.rowCount())
+                    if class_table.item(row, 1) is not None
+                ]
+
+                if not ordered_ids:
+                    self.show_notification("No classes defined for mask saving.")
+                else:
+                    id_map = {
+                        old_id: new_id for new_id, old_id in enumerate(ordered_ids)
+                    }
+                    num_final_classes = len(ordered_ids)
+                    final_mask_tensor = np.zeros(
+                        (h, w, num_final_classes), dtype=np.uint8
+                    )
 
-        np.savez_compressed(output_path, mask=final_mask_tensor.astype(np.uint8))
+                    for seg in self.segments:
+                        class_id = seg.get("class_id")
+                        if class_id not in id_map:
+                            continue
+                        new_channel_idx = id_map[class_id]
+                        mask = (
+                            self.rasterize_polygon(seg["vertices"])
+                            if seg["type"] == "Polygon"
+                            else seg.get("mask")
+                        )
+                        if mask is not None:
+                            final_mask_tensor[:, :, new_channel_idx] = np.logical_or(
+                                final_mask_tensor[:, :, new_channel_idx], mask
+                            )
+                    if save_npz:
+                        npz_path = os.path.splitext(self.current_image_path)[0] + ".npz"
+                        np.savez_compressed(
+                            npz_path, mask=final_mask_tensor.astype(np.uint8)
+                        )
+                        self.file_model.update_cache_for_path(npz_path)
 
-        self.file_model.set_highlighted_path(output_path)
-        QTimer.singleShot(1500, lambda: self.file_model.set_highlighted_path(None))
+                        self.file_model.set_highlighted_path(npz_path)
+                        QTimer.singleShot(
+                            1500, lambda: self.file_model.set_highlighted_path(None)
+                        )
+                        saved_something = True
+                    if save_txt:
+                        if self.control_panel.chk_yolo_use_alias.isChecked():
+                            class_labels = [
+                                class_table.item(row, 0).text()
+                                for row in range(class_table.rowCount())
+                            ]
+                        else:
+                            class_labels = list(range(num_final_classes))
+
+                        txt_path = self.generate_yolo_annotations(
+                            final_mask_tensor, class_labels
+                        )
+                        if txt_path:
+                            self.file_model.update_cache_for_path(txt_path)
+                        saved_something = True
+
+        if save_aliases:
+            aliases_path = os.path.splitext(self.current_image_path)[0] + ".json"
+            aliases_to_save = {str(k): v for k, v in self.class_aliases.items()}
+            with open(aliases_path, "w") as f:
+                json.dump(aliases_to_save, f, indent=4)
+            saved_something = True
+
+        if saved_something:
+            self.right_panel.status_label.setText("Saved!")
+        else:
+            self.right_panel.status_label.clear()
 
-        self.right_panel.status_label.setText("Saved!")
-        self.generate_yolo_annotations(npz_file_path=output_path)
         QTimer.singleShot(3000, lambda: self.right_panel.status_label.clear())
 
-    def generate_yolo_annotations(self, npz_file_path):
+    def generate_yolo_annotations(self, mask_tensor, class_labels):
         output_path = os.path.splitext(self.current_image_path)[0] + ".txt"
-        npz_data = np.load(npz_file_path)  # Load the saved npz file
-
-        img = npz_data["mask"][:, :, :]
-        num_channels = img.shape[2]  # C
-        h, w = img.shape[:2]  # H, W
+        h, w, num_channels = mask_tensor.shape
 
         directory_path = os.path.dirname(output_path)
         os.makedirs(directory_path, exist_ok=True)
 
         yolo_annotations = []
-
         for channel in range(num_channels):
-            single_channel_image = img[:, :, channel]
+            single_channel_image = mask_tensor[:, :, channel]
+            if not np.any(single_channel_image):
+                continue
+
             contours, _ = cv2.findContours(
                 single_channel_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
             )
 
-            class_id = channel  # Use the channel index as the class ID
-
+            class_label = class_labels[channel]
             for contour in contours:
                 x, y, width, height = cv2.boundingRect(contour)
-                center_x = x + width / 2
-                center_y = y + height / 2
-
-                normalized_center_x = center_x / w
-                normalized_center_y = center_y / h
+                center_x = (x + width / 2) / w
+                center_y = (y + height / 2) / h
                 normalized_width = width / w
                 normalized_height = height / h
-
-                yolo_entry = f"{class_id} {normalized_center_x} {normalized_center_y} {normalized_width} {normalized_height}"
+                yolo_entry = f"{class_label} {center_x} {center_y} {normalized_width} {normalized_height}"
                 yolo_annotations.append(yolo_entry)
 
+        if not yolo_annotations:
+            return None
+
         with open(output_path, "w") as file:
             for annotation in yolo_annotations:
                 file.write(annotation + "\n")
 
+        return output_path
+
     def save_current_segment(self):
         if (
             self.mode != "sam_points"
@@ -1040,6 +1120,20 @@ class MainWindow(QMainWindow):
         self.update_all_lists()
         self.viewer.setFocus()
 
+    def load_class_aliases(self):
+        if not self.current_image_path:
+            return
+        json_path = os.path.splitext(self.current_image_path)[0] + ".json"
+        if os.path.exists(json_path):
+            try:
+                with open(json_path, "r") as f:
+                    loaded_aliases = json.load(f)
+                    # JSON loads keys as strings, convert them to int
+                    self.class_aliases = {int(k): v for k, v in loaded_aliases.items()}
+            except (json.JSONDecodeError, ValueError) as e:
+                print(f"Error loading class aliases from {json_path}: {e}")
+                self.class_aliases.clear()
+
     def load_existing_mask(self):
         if not self.current_image_path:
             return
@@ -1068,12 +1162,10 @@ class MainWindow(QMainWindow):
     def add_point(self, pos, positive):
         point_list = self.positive_points if positive else self.negative_points
         point_list.append([int(pos.x()), int(pos.y())])
-
         point_color = (
             QColor(Qt.GlobalColor.green) if positive else QColor(Qt.GlobalColor.red)
         )
         point_color.setAlpha(150)
-
         point_diameter = self.point_radius * 2
         point_item = QGraphicsEllipseItem(
             pos.x() - self.point_radius,
@@ -1120,17 +1212,15 @@ class MainWindow(QMainWindow):
                 (pos.x() - self.polygon_points[0].x()) ** 2
                 + (pos.y() - self.polygon_points[0].y()) ** 2
             )
-            < 4  # pixel distance threshold squared
+            < self.polygon_join_threshold**2
         ):
             if len(self.polygon_points) > 2:
                 self.finalize_polygon()
             return
         self.polygon_points.append(pos)
         point_diameter = self.point_radius * 2
-
         point_color = QColor(Qt.GlobalColor.blue)
         point_color.setAlpha(150)
-
         dot = QGraphicsEllipseItem(
             pos.x() - self.point_radius,
             pos.y() - self.point_radius,
@@ -1144,7 +1234,6 @@ class MainWindow(QMainWindow):
         self.draw_polygon_preview()
 
     def draw_polygon_preview(self):
-        # Clean up old preview lines/polygons
         for item in self.polygon_preview_items:
             if not isinstance(item, QGraphicsEllipseItem):
                 if item.scene():
@@ -1154,7 +1243,6 @@ class MainWindow(QMainWindow):
             for item in self.polygon_preview_items
             if isinstance(item, QGraphicsEllipseItem)
         ]
-
         if len(self.polygon_points) > 2:
             preview_poly = QGraphicsPolygonItem(QPolygonF(self.polygon_points))
             preview_poly.setBrush(QBrush(QColor(0, 255, 255, 100)))