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

lazylabel/main.py

@@ -16,8 +16,18 @@ from PyQt6.QtWidgets import (
     QGraphicsPolygonItem,
     QTableWidgetSelectionRange,
 )
-from PyQt6.QtGui import QPixmap, QColor, QPen, QBrush, QPolygonF, QIcon
-from PyQt6.QtCore import Qt, QPointF, QTimer
+from PyQt6.QtGui import (
+    QPixmap,
+    QColor,
+    QPen,
+    QBrush,
+    QPolygonF,
+    QIcon,
+    QCursor,
+    QKeySequence,
+    QShortcut,
+)
+from PyQt6.QtCore import Qt, QPointF, QTimer, QModelIndex
 
 from .photo_viewer import PhotoViewer
 from .sam_model import SamModel
@@ -26,6 +36,7 @@ from .controls import ControlPanel, RightPanel
 from .custom_file_system_model import CustomFileSystemModel
 from .editable_vertex import EditableVertexItem
 from .hoverable_polygon_item import HoverablePolygonItem
+from .hoverable_pixelmap_item import HoverablePixmapItem
 from .numeric_table_widget_item import NumericTableWidgetItem
 
 
@@ -34,7 +45,9 @@ class MainWindow(QMainWindow):
         super().__init__()
         self.setWindowTitle("LazyLabel by DNC")
 
-        icon_path = os.path.join(os.path.dirname(__file__), "demo_pictures", "logo2.png")
+        icon_path = os.path.join(
+            os.path.dirname(__file__), "demo_pictures", "logo2.png"
+        )
         if os.path.exists(icon_path):
             self.setWindowIcon(QIcon(icon_path))
 
@@ -44,9 +57,18 @@ class MainWindow(QMainWindow):
         self.mode = "sam_points"
         self.previous_mode = "sam_points"
         self.current_image_path = None
-        self.current_file_index = None
+        self.current_file_index = QModelIndex()
+
         self.next_class_id = 0
 
+        self.class_aliases = {}  # {class_id: "alias_string"}
+
+        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.point_items, self.positive_points, self.negative_points = [], [], []
         self.polygon_points, self.polygon_preview_items = [], []
         self.rubber_band_line = None
@@ -74,7 +96,9 @@ class MainWindow(QMainWindow):
         central_widget.setLayout(main_layout)
         self.setCentralWidget(central_widget)
 
-        self.control_panel.device_label.setText(f"Device: {str(self.sam_model.device).upper()}")
+        self.control_panel.device_label.setText(
+            f"Device: {str(self.sam_model.device).upper()}"
+        )
         self.setup_connections()
         self.set_sam_mode()
 
@@ -89,24 +113,48 @@ class MainWindow(QMainWindow):
 
         self.right_panel.btn_open_folder.clicked.connect(self.open_folder_dialog)
         self.right_panel.file_tree.doubleClicked.connect(self.load_selected_image)
-        self.right_panel.btn_merge_selection.clicked.connect(self.assign_selected_to_class)
-        self.right_panel.btn_delete_selection.clicked.connect(self.delete_selected_segments)
-        self.right_panel.segment_table.itemSelectionChanged.connect(self.highlight_selected_segments)
+        self.right_panel.btn_merge_selection.clicked.connect(
+            self.assign_selected_to_class
+        )
+        self.right_panel.btn_delete_selection.clicked.connect(
+            self.delete_selected_segments
+        )
+        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(self.update_segment_table)
+        self.right_panel.class_filter_combo.currentIndexChanged.connect(
+            self.update_segment_table
+        )
 
         self.control_panel.btn_sam_mode.clicked.connect(self.set_sam_mode)
         self.control_panel.btn_polygon_mode.clicked.connect(self.set_polygon_mode)
-        self.control_panel.btn_selection_mode.clicked.connect(self.toggle_selection_mode)
+        self.control_panel.btn_selection_mode.clicked.connect(
+            self.toggle_selection_mode
+        )
         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)
+
+    def show_notification(self, message, duration=3000):
+        self.control_panel.notification_label.setText(message)
+        QTimer.singleShot(
+            duration, lambda: self.control_panel.notification_label.clear()
+        )
 
-    def _get_color_for_class(self, class_id, saturation, value):
+    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, saturation, value)
+        color = QColor.fromHsv(hue, 220, 220)
 
         if not color.isValid():
             return QColor(Qt.GlobalColor.white)
@@ -122,9 +170,25 @@ class MainWindow(QMainWindow):
             self.previous_mode = self.mode
 
         self.mode = mode_name
-        self.control_panel.mode_label.setText(f"Mode: {mode_name.replace('_', ' ').title()}")
+        self.control_panel.mode_label.setText(
+            f"Mode: {mode_name.replace('_', ' ').title()}"
+        )
         self.clear_all_points()
-        self.viewer.setDragMode(self.viewer.DragMode.ScrollHandDrag if self.mode == "pan" else self.viewer.DragMode.NoDrag)
+
+        cursor_map = {
+            "sam_points": Qt.CursorShape.CrossCursor,
+            "polygon": Qt.CursorShape.CrossCursor,
+            "selection": Qt.CursorShape.ArrowCursor,
+            "edit": Qt.CursorShape.SizeAllCursor,
+            "pan": Qt.CursorShape.OpenHandCursor,
+        }
+        self.viewer.set_cursor(cursor_map.get(self.mode, Qt.CursorShape.ArrowCursor))
+
+        self.viewer.setDragMode(
+            self.viewer.DragMode.ScrollHandDrag
+            if self.mode == "pan"
+            else self.viewer.DragMode.NoDrag
+        )
 
     def set_sam_mode(self):
         self.set_mode("sam_points")
@@ -148,27 +212,44 @@ class MainWindow(QMainWindow):
 
     def toggle_edit_mode(self):
         selected_indices = self.get_selected_segment_indices()
-        can_edit = any(self.segments[i].get("type") == "Polygon" for i in selected_indices)
+
         if self.mode == "edit":
             self.set_mode("selection", is_toggle=True)
-        elif self.mode == "selection" and can_edit:
-            self.set_mode("edit", is_toggle=True)
-            self.display_all_segments()
+            return
+
+        if not selected_indices:
+            self.show_notification("Select a polygon to edit.")
+            return
+
+        can_edit = any(
+            self.segments[i].get("type") == "Polygon" for i in selected_indices
+        )
+
+        if not can_edit:
+            self.show_notification("Only polygon segments can be edited.")
+            return
+
+        self.set_mode("edit", is_toggle=True)
+        self.display_all_segments()
 
     def open_folder_dialog(self):
         folder_path = QFileDialog.getExistingDirectory(self, "Select Image Folder")
         if folder_path:
-            self.right_panel.file_tree.setRootIndex(self.file_model.setRootPath(folder_path))
+            self.right_panel.file_tree.setRootIndex(
+                self.file_model.setRootPath(folder_path)
+            )
         self.viewer.setFocus()
 
     def load_selected_image(self, index):
-        if not index.isValid():
+        if not index.isValid() or not self.file_model.isDir(index.parent()):
             return
 
         self.current_file_index = index
         path = self.file_model.filePath(index)
 
-        if os.path.isfile(path) and path.lower().endswith((".png", ".jpg", ".jpeg", ".tiff", ".tif")):
+        if os.path.isfile(path) and path.lower().endswith(
+            (".png", ".jpg", ".jpeg", ".tiff", ".tif")
+        ):
             self.current_image_path = path
             pixmap = QPixmap(self.current_image_path)
             if not pixmap.isNull():
@@ -176,14 +257,47 @@ class MainWindow(QMainWindow):
                 self.viewer.set_photo(pixmap)
                 self.sam_model.set_image(self.current_image_path)
                 self.load_existing_mask()
+                self.right_panel.file_tree.setCurrentIndex(index)
                 self.viewer.setFocus()
 
+    def load_next_image(self):
+        if not self.current_file_index.isValid():
+            return
+
+        if self.control_panel.chk_auto_save.isChecked():
+            self.save_output_to_npz()
+
+        row = self.current_file_index.row()
+        parent = self.current_file_index.parent()
+        if row + 1 < self.file_model.rowCount(parent):
+            next_index = self.file_model.index(row + 1, 0, parent)
+            self.load_selected_image(next_index)
+
+    def load_previous_image(self):
+        if not self.current_file_index.isValid():
+            return
+
+        if self.control_panel.chk_auto_save.isChecked():
+            self.save_output_to_npz()
+
+        row = self.current_file_index.row()
+        parent = self.current_file_index.parent()
+        if row > 0:
+            prev_index = self.file_model.index(row - 1, 0, parent)
+            self.load_selected_image(prev_index)
+
     def reset_state(self):
         self.clear_all_points()
+        # Preserve aliases between images in the same session
+        # self.class_aliases.clear()
         self.segments.clear()
         self.next_class_id = 0
         self.update_all_lists()
-        items_to_remove = [item for item in self.viewer.scene().items() if item is not self.viewer._pixmap_item]
+        items_to_remove = [
+            item
+            for item in self.viewer.scene().items()
+            if item is not self.viewer._pixmap_item
+        ]
         for item in items_to_remove:
             self.viewer.scene().removeItem(item)
         self.segment_items.clear()
@@ -193,14 +307,32 @@ class MainWindow(QMainWindow):
         key, mods = event.key(), event.modifiers()
         if event.isAutoRepeat():
             return
+
+        pan_multiplier = 5.0 if (mods & Qt.KeyboardModifier.ShiftModifier) else 2
+
         if key == Qt.Key.Key_W:
-            self.viewer.verticalScrollBar().setValue(self.viewer.verticalScrollBar().value() - int(self.viewer.height() * 0.1))
-        elif key == Qt.Key.Key_S and not mods:
-            self.viewer.verticalScrollBar().setValue(self.viewer.verticalScrollBar().value() + int(self.viewer.height() * 0.1))
+            amount = int(self.viewer.height() * 0.1 * pan_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)
+            self.viewer.verticalScrollBar().setValue(
+                self.viewer.verticalScrollBar().value() + amount
+            )
         elif key == Qt.Key.Key_A and not (mods & Qt.KeyboardModifier.ControlModifier):
-            self.viewer.horizontalScrollBar().setValue(self.viewer.horizontalScrollBar().value() - int(self.viewer.width() * 0.1))
+            amount = int(self.viewer.width() * 0.1 * pan_multiplier)
+            self.viewer.horizontalScrollBar().setValue(
+                self.viewer.horizontalScrollBar().value() - amount
+            )
         elif key == Qt.Key.Key_D:
-            self.viewer.horizontalScrollBar().setValue(self.viewer.horizontalScrollBar().value() + int(self.viewer.width() * 0.1))
+            amount = int(self.viewer.width() * 0.1 * pan_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:
@@ -223,16 +355,47 @@ class MainWindow(QMainWindow):
         elif key == Qt.Key.Key_A and mods == Qt.KeyboardModifier.ControlModifier:
             self.right_panel.segment_table.selectAll()
         elif key == Qt.Key.Key_Space:
-            self.save_current_segment()
+            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:
-            self.save_output_to_npz()
+            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()
+        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()
 
     def scene_mouse_press(self, event):
         self._original_mouse_press(event)
         if event.isAccepted():
             return
+
+        if self.mode == "pan":
+            self.viewer.set_cursor(Qt.CursorShape.ClosedHandCursor)
+
         pos = event.scenePos()
-        if self.viewer._pixmap_item.pixmap().isNull() or not self.viewer._pixmap_item.pixmap().rect().contains(pos.toPoint()):
+        if (
+            self.viewer._pixmap_item.pixmap().isNull()
+            or not self.viewer._pixmap_item.pixmap().rect().contains(pos.toPoint())
+        ):
             return
         if self.mode == "sam_points":
             if event.button() == Qt.MouseButton.LeftButton:
@@ -252,7 +415,9 @@ class MainWindow(QMainWindow):
             self.is_dragging_polygon = True
             selected_indices = self.get_selected_segment_indices()
             self.drag_initial_vertices = {
-                i: list(self.segments[i]["vertices"]) for i in selected_indices if self.segments[i].get("type") == "Polygon"
+                i: list(self.segments[i]["vertices"])
+                for i in selected_indices
+                if self.segments[i].get("type") == "Polygon"
             }
 
     def scene_mouse_move(self, event):
@@ -260,13 +425,22 @@ class MainWindow(QMainWindow):
         if self.mode == "edit" and self.is_dragging_polygon:
             delta = pos - self.drag_start_pos
             for i, initial_verts in self.drag_initial_vertices.items():
-                self.segments[i]["vertices"] = [QPointF(v.x() + delta.x(), v.y() + delta.y()) for v in initial_verts]
+                self.segments[i]["vertices"] = [
+                    QPointF(v.x() + delta.x(), v.y() + delta.y()) for v in initial_verts
+                ]
                 self.update_polygon_visuals(i)
         elif self.mode == "polygon" and self.polygon_points:
             if self.rubber_band_line is None:
                 self.rubber_band_line = QGraphicsLineItem()
-                self.rubber_band_line.setPen(QPen(Qt.GlobalColor.white, 2, Qt.PenStyle.DotLine))
+
+                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(),
@@ -278,6 +452,9 @@ class MainWindow(QMainWindow):
             self._original_mouse_move(event)
 
     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()
@@ -286,20 +463,55 @@ class MainWindow(QMainWindow):
     def undo_last_action(self):
         if self.mode == "polygon" and self.polygon_points:
             self.polygon_points.pop()
-            if self.polygon_preview_items:
-                self.viewer.scene().removeItem(self.polygon_preview_items.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)
+                point_color.setAlpha(150)
+                dot = QGraphicsEllipseItem(
+                    point.x() - self.point_radius,
+                    point.y() - self.point_radius,
+                    point_diameter,
+                    point_diameter,
+                )
+                dot.setBrush(QBrush(point_color))
+                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(4, 4)
+            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()
 
+    def _update_next_class_id(self):
+        all_ids = {
+            seg.get("class_id")
+            for seg in self.segments
+            if seg.get("class_id") is not None
+        }
+        if not all_ids:
+            self.next_class_id = 0
+        else:
+            self.next_class_id = max(all_ids) + 1
+
     def finalize_polygon(self):
         if len(self.polygon_points) < 3:
             return
@@ -314,7 +526,7 @@ class MainWindow(QMainWindow):
                 "class_id": self.next_class_id,
             }
         )
-        self.next_class_id += 1
+        self._update_next_class_id()
         self.polygon_points.clear()
         for item in self.polygon_preview_items:
             self.viewer.scene().removeItem(item)
@@ -325,14 +537,25 @@ class MainWindow(QMainWindow):
         x, y = int(pos.x()), int(pos.y())
         for i in range(len(self.segments) - 1, -1, -1):
             seg = self.segments[i]
-            mask = self.rasterize_polygon(seg["vertices"]) if seg["type"] == "Polygon" else seg.get("mask")
-            if mask is not None and y < mask.shape[0] and x < mask.shape[1] and mask[y, x]:
+            mask = (
+                self.rasterize_polygon(seg["vertices"])
+                if seg["type"] == "Polygon"
+                else seg.get("mask")
+            )
+            if (
+                mask is not None
+                and y < mask.shape[0]
+                and x < mask.shape[1]
+                and mask[y, x]
+            ):
                 for j in range(self.right_panel.segment_table.rowCount()):
                     item = self.right_panel.segment_table.item(j, 0)
                     if item and item.data(Qt.ItemDataRole.UserRole) == i:
                         table = self.right_panel.segment_table
                         is_selected = table.item(j, 0).isSelected()
-                        range_to_select = QTableWidgetSelectionRange(j, 0, j, table.columnCount() - 1)
+                        range_to_select = QTableWidgetSelectionRange(
+                            j, 0, j, table.columnCount() - 1
+                        )
                         table.setRangeSelected(range_to_select, not is_selected)
                         return
         self.viewer.setFocus()
@@ -342,7 +565,11 @@ class MainWindow(QMainWindow):
         if not selected_indices:
             return
 
-        existing_class_ids = [self.segments[i]["class_id"] for i in selected_indices if self.segments[i].get("class_id") is not None]
+        existing_class_ids = [
+            self.segments[i]["class_id"]
+            for i in selected_indices
+            if self.segments[i].get("class_id") is not None
+        ]
 
         if existing_class_ids:
             target_class_id = min(existing_class_ids)
@@ -352,6 +579,7 @@ class MainWindow(QMainWindow):
         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()
@@ -378,31 +606,62 @@ 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, saturation=220, value=220)
+
+            base_color = self._get_color_for_class(class_id)
 
             if seg_dict["type"] == "Polygon":
                 poly_item = HoverablePolygonItem(QPolygonF(seg_dict["vertices"]))
-                default_brush = QBrush(QColor(base_color.red(), base_color.green(), base_color.blue(), 70))
-                hover_brush = QBrush(QColor(base_color.red(), base_color.green(), base_color.blue(), 170))
+                default_brush = QBrush(
+                    QColor(base_color.red(), base_color.green(), base_color.blue(), 70)
+                )
+                hover_brush = QBrush(
+                    QColor(base_color.red(), base_color.green(), base_color.blue(), 170)
+                )
                 poly_item.set_brushes(default_brush, hover_brush)
                 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
                 for v in seg_dict["vertices"]:
-                    dot = QGraphicsEllipseItem(v.x() - 3, v.y() - 3, 6, 6)
+                    dot = QGraphicsEllipseItem(
+                        v.x() - self.point_radius,
+                        v.y() - self.point_radius,
+                        point_diameter,
+                        point_diameter,
+                    )
                     dot.setBrush(vertex_color)
+                    dot.setPen(QPen(Qt.GlobalColor.transparent))
                     self.viewer.scene().addItem(dot)
                     self.segment_items[i].append(dot)
                 if self.mode == "edit" and i in selected_indices:
+                    handle_diameter = self.point_radius * 2
                     for idx, v in enumerate(seg_dict["vertices"]):
-                        vertex_item = EditableVertexItem(self, i, idx, -4, -4, 8, 8)
+                        vertex_item = EditableVertexItem(
+                            self,
+                            i,
+                            idx,
+                            -handle_diameter / 2,
+                            -handle_diameter / 2,
+                            handle_diameter,
+                            handle_diameter,
+                        )
                         vertex_item.setPos(v)
                         self.viewer.scene().addItem(vertex_item)
                         self.segment_items[i].append(vertex_item)
             elif seg_dict.get("mask") is not None:
-                pixmap = mask_to_pixmap(seg_dict["mask"], base_color.getRgb()[:3])
-                pixmap_item = self.viewer.scene().addPixmap(pixmap)
+                default_pixmap = mask_to_pixmap(
+                    seg_dict["mask"], base_color.getRgb()[:3], alpha=70
+                )
+                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)
                 pixmap_item.setZValue(i + 1)
                 self.segment_items[i].append(pixmap_item)
         self.highlight_selected_segments()
@@ -426,7 +685,11 @@ class MainWindow(QMainWindow):
         selected_indices = self.get_selected_segment_indices()
         for i in selected_indices:
             seg = self.segments[i]
-            mask = self.rasterize_polygon(seg["vertices"]) if seg["type"] == "Polygon" else seg.get("mask")
+            mask = (
+                self.rasterize_polygon(seg["vertices"])
+                if seg["type"] == "Polygon"
+                else seg.get("mask")
+            )
             if mask is not None:
                 pixmap = mask_to_pixmap(mask, (255, 255, 255))
                 highlight_item = self.viewer.scene().addPixmap(pixmap)
@@ -434,9 +697,9 @@ 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_filter_combo()
         self.update_segment_table()
-        self.update_class_list()
         self.display_all_segments()
 
     def update_segment_table(self):
@@ -450,7 +713,7 @@ class MainWindow(QMainWindow):
         filter_class_id = -1
         if not show_all:
             try:
-                filter_class_id = int(filter_text.split(" ")[1])
+                filter_class_id = int(filter_text.split("(ID: ")[1][:-1])
             except (ValueError, IndexError):
                 pass
 
@@ -463,7 +726,7 @@ 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, saturation=180, value=200)
+            color = self._get_color_for_class(class_id)
 
             class_id_str = str(class_id) if class_id is not None else "N/A"
             index_item = NumericTableWidgetItem(str(original_index + 1))
@@ -478,7 +741,7 @@ class MainWindow(QMainWindow):
             table.setItem(row, 1, class_item)
             table.setItem(row, 2, type_item)
 
-            for col in range(3):
+            for col in range(table.columnCount()):
                 if table.item(row, col):
                     table.item(row, col).setBackground(QBrush(color))
 
@@ -495,30 +758,70 @@ class MainWindow(QMainWindow):
     def update_class_list(self):
         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()
 
-        unique_class_ids = sorted(list({seg.get("class_id") for seg in self.segments if seg.get("class_id") is not None}))
+        unique_class_ids = sorted(
+            list(
+                {
+                    seg.get("class_id")
+                    for seg in self.segments
+                    if seg.get("class_id") is not None
+                }
+            )
+        )
         class_table.setRowCount(len(unique_class_ids))
 
         for row, cid in enumerate(unique_class_ids):
-            item = QTableWidgetItem(str(cid))
-            item.setFlags(item.flags() & ~Qt.ItemFlag.ItemIsEditable)
+            alias = self.class_aliases.get(cid, str(cid))
+            alias_item = QTableWidgetItem(alias)
+            id_item = QTableWidgetItem(str(cid))
 
-            color = self._get_color_for_class(cid, saturation=180, value=200)
+            id_item.setFlags(id_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
 
-            item.setBackground(QBrush(color))
-            class_table.setItem(row, 0, item)
+            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)
 
         class_table.blockSignals(False)
 
     def update_class_filter_combo(self):
         combo = self.right_panel.class_filter_combo
-        unique_class_ids = sorted(list({seg.get("class_id") for seg in self.segments if seg.get("class_id") is not None}))
+        unique_class_ids = sorted(
+            list(
+                {
+                    seg.get("class_id")
+                    for seg in self.segments
+                    if seg.get("class_id") is not None
+                }
+            )
+        )
+
         current_selection = combo.currentText()
         combo.blockSignals(True)
         combo.clear()
         combo.addItem("All Classes")
-        combo.addItems([f"Class {cid}" for cid in unique_class_ids])
+        combo.addItems(
+            [
+                f"{self.class_aliases.get(cid, cid)} (ID: {cid})"
+                for cid in unique_class_ids
+            ]
+        )
+
         if combo.findText(current_selection) > -1:
             combo.setCurrentText(current_selection)
         else:
@@ -527,20 +830,54 @@ class MainWindow(QMainWindow):
 
     def reassign_class_ids(self):
         class_table = self.right_panel.class_table
-        ordered_ids = [
-            int(class_table.item(row, 0).text()) for row in range(class_table.rowCount()) if class_table.item(row, 0) is not None
-        ]
+
+        ordered_ids = []
+        for row in range(class_table.rowCount()):
+            id_item = class_table.item(row, 1)
+            if id_item and id_item.text():
+                try:
+                    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]
-        self.next_class_id = len(ordered_ids)
+
+        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
+            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
+
+        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:
+        if item.column() != 1:  # Class ID column in segment table
             return
         table = self.right_panel.segment_table
         index_item = table.item(item.row(), 0)
@@ -559,14 +896,15 @@ class MainWindow(QMainWindow):
                 raise IndexError("Invalid segment index found in table.")
 
             self.segments[original_index]["class_id"] = new_class_id
-            if new_class_id >= self.next_class_id:
-                self.next_class_id = new_class_id + 1
+            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")
+                item.setText(
+                    str(original_class_id) if original_class_id is not None else "N/A"
+                )
         finally:
             table.blockSignals(False)
             self.viewer.setFocus()
@@ -575,7 +913,11 @@ class MainWindow(QMainWindow):
         table = self.right_panel.segment_table
         selected_items = table.selectedItems()
         selected_rows = sorted(list({item.row() for item in selected_items}))
-        return [table.item(row, 0).data(Qt.ItemDataRole.UserRole) for row in selected_rows if table.item(row, 0)]
+        return [
+            table.item(row, 0).data(Qt.ItemDataRole.UserRole)
+            for row in selected_rows
+            if table.item(row, 0)
+        ]
 
     def save_output_to_npz(self):
         if not self.segments or not self.current_image_path:
@@ -588,14 +930,21 @@ class MainWindow(QMainWindow):
             self.viewer._pixmap_item.pixmap().height(),
             self.viewer._pixmap_item.pixmap().width(),
         )
-        unique_class_ids = sorted(list({seg["class_id"] for seg in self.segments if seg.get("class_id") is not None}))
-        if not unique_class_ids:
-            self.right_panel.status_label.setText("Save failed: No classes.")
+
+        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(unique_class_ids)}
-        num_final_classes = len(unique_class_ids)
+        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)
 
         for seg in self.segments:
@@ -603,12 +952,20 @@ class MainWindow(QMainWindow):
             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")
+            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)
+                final_mask_tensor[:, :, new_channel_idx] = np.logical_or(
+                    final_mask_tensor[:, :, new_channel_idx], mask
+                )
 
         np.savez_compressed(output_path, mask=final_mask_tensor.astype(np.uint8))
-        self.file_model.setRootPath(self.file_model.rootPath())
+
+        self.file_model.set_highlighted_path(output_path)
+        QTimer.singleShot(1500, lambda: self.file_model.set_highlighted_path(None))
 
         self.right_panel.status_label.setText("Saved!")
         self.generate_yolo_annotations(npz_file_path=output_path)
@@ -629,7 +986,9 @@ class MainWindow(QMainWindow):
 
         for channel in range(num_channels):
             single_channel_image = img[:, :, channel]
-            contours, _ = cv2.findContours(single_channel_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            contours, _ = cv2.findContours(
+                single_channel_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+            )
 
             class_id = channel  # Use the channel index as the class ID
 
@@ -651,7 +1010,11 @@ class MainWindow(QMainWindow):
                 file.write(annotation + "\n")
 
     def save_current_segment(self):
-        if self.mode != "sam_points" or not hasattr(self, "preview_mask_item") or not self.preview_mask_item:
+        if (
+            self.mode != "sam_points"
+            or not hasattr(self, "preview_mask_item")
+            or not self.preview_mask_item
+        ):
             return
         mask = self.sam_model.predict(self.positive_points, self.negative_points)
         if mask is not None:
@@ -663,7 +1026,7 @@ class MainWindow(QMainWindow):
                     "class_id": self.next_class_id,
                 }
             )
-            self.next_class_id += 1
+            self._update_next_class_id()
             self.clear_all_points()
             self.update_all_lists()
 
@@ -673,6 +1036,7 @@ class MainWindow(QMainWindow):
             return
         for i in sorted(selected_indices, reverse=True):
             del self.segments[i]
+        self._update_next_class_id()
         self.update_all_lists()
         self.viewer.setFocus()
 
@@ -698,16 +1062,27 @@ class MainWindow(QMainWindow):
                                     "class_id": i,
                                 }
                             )
-                    self.next_class_id = num_classes
+                    self._update_next_class_id()
             self.update_all_lists()
 
     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())])
-        color = Qt.GlobalColor.green if positive else Qt.GlobalColor.red
-        point_item = QGraphicsEllipseItem(pos.x() - 4, pos.y() - 4, 8, 8)
-        point_item.setBrush(QBrush(color))
-        point_item.setPen(QPen(Qt.GlobalColor.white))
+
+        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,
+            pos.y() - self.point_radius,
+            point_diameter,
+            point_diameter,
+        )
+        point_item.setBrush(QBrush(point_color))
+        point_item.setPen(QPen(Qt.GlobalColor.transparent))
         self.viewer.scene().addItem(point_item)
         self.point_items.append(point_item)
 
@@ -740,26 +1115,45 @@ class MainWindow(QMainWindow):
             self.preview_mask_item = None
 
     def handle_polygon_click(self, pos):
-        if self.polygon_points and (((pos.x() - self.polygon_points[0].x()) ** 2 + (pos.y() - self.polygon_points[0].y()) ** 2) < 25):
+        if self.polygon_points and (
+            (
+                (pos.x() - self.polygon_points[0].x()) ** 2
+                + (pos.y() - self.polygon_points[0].y()) ** 2
+            )
+            < 4  # pixel distance threshold squared
+        ):
             if len(self.polygon_points) > 2:
                 self.finalize_polygon()
             return
         self.polygon_points.append(pos)
-        dot = QGraphicsEllipseItem(pos.x() - 2, pos.y() - 2, 4, 4)
-        dot.setBrush(QBrush(Qt.GlobalColor.blue))
-        dot.setPen(QPen(Qt.GlobalColor.cyan))
+        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,
+            point_diameter,
+            point_diameter,
+        )
+        dot.setBrush(QBrush(point_color))
+        dot.setPen(QPen(Qt.GlobalColor.transparent))
         self.viewer.scene().addItem(dot)
         self.polygon_preview_items.append(dot)
         self.draw_polygon_preview()
 
     def draw_polygon_preview(self):
-        if self.rubber_band_line:
-            self.viewer.scene().removeItem(self.rubber_band_line)
-            self.rubber_band_line = None
+        # Clean up old preview lines/polygons
         for item in self.polygon_preview_items:
             if not isinstance(item, QGraphicsEllipseItem):
-                self.viewer.scene().removeItem(item)
-        self.polygon_preview_items = [item for item in self.polygon_preview_items if isinstance(item, QGraphicsEllipseItem)]
+                if item.scene():
+                    self.viewer.scene().removeItem(item)
+        self.polygon_preview_items = [
+            item
+            for item in self.polygon_preview_items
+            if isinstance(item, QGraphicsEllipseItem)
+        ]
 
         if len(self.polygon_points) > 2:
             preview_poly = QGraphicsPolygonItem(QPolygonF(self.polygon_points))
@@ -769,6 +1163,8 @@ class MainWindow(QMainWindow):
             self.polygon_preview_items.append(preview_poly)
 
         if len(self.polygon_points) > 1:
+            line_color = QColor(Qt.GlobalColor.cyan)
+            line_color.setAlpha(150)
             for i in range(len(self.polygon_points) - 1):
                 line = QGraphicsLineItem(
                     self.polygon_points[i].x(),
@@ -776,7 +1172,7 @@ class MainWindow(QMainWindow):
                     self.polygon_points[i + 1].x(),
                     self.polygon_points[i + 1].y(),
                 )
-                line.setPen(QPen(Qt.GlobalColor.cyan, 2))
+                line.setPen(QPen(line_color, self.line_thickness))
                 self.viewer.scene().addItem(line)
                 self.polygon_preview_items.append(line)