Implementing 2 stage training

src/gui/widgets/annotation_canvas_widget.py

@@ -17,7 +17,7 @@ from PySide6.QtGui import (
     QMouseEvent,
     QPaintEvent,
 )
-from PySide6.QtCore import Qt, QEvent, Signal, QPoint
+from PySide6.QtCore import Qt, QEvent, Signal, QPoint, QRect
 from typing import Any, Dict, List, Optional, Tuple
 
 from src.utils.image import Image, ImageLoadError
@@ -529,6 +529,40 @@ class AnnotationCanvasWidget(QWidget):
                 painter.setPen(pen)
                 painter.drawRect(x_min, y_min, rect_width, rect_height)
 
+                label_text = meta.get("label")
+                if label_text:
+                    painter.save()
+                    font = painter.font()
+                    font.setPointSizeF(max(10.0, width + 4))
+                    painter.setFont(font)
+                    metrics = painter.fontMetrics()
+                    text_width = metrics.horizontalAdvance(label_text)
+                    text_height = metrics.height()
+                    padding = 4
+                    bg_width = text_width + padding * 2
+                    bg_height = text_height + padding * 2
+                    canvas_width = self.original_pixmap.width()
+                    canvas_height = self.original_pixmap.height()
+                    bg_x = max(0, min(x_min, canvas_width - bg_width))
+                    bg_y = y_min - bg_height
+                    if bg_y < 0:
+                        bg_y = min(y_min, canvas_height - bg_height)
+                        bg_y = max(0, bg_y)
+                    background_rect = QRect(bg_x, bg_y, bg_width, bg_height)
+                    background_color = QColor(pen_color)
+                    background_color.setAlpha(220)
+                    painter.fillRect(background_rect, background_color)
+                    text_color = QColor(0, 0, 0)
+                    if background_color.lightness() < 128:
+                        text_color = QColor(255, 255, 255)
+                    painter.setPen(text_color)
+                    painter.drawText(
+                        background_rect.adjusted(padding, padding, -padding, -padding),
+                        Qt.AlignLeft | Qt.AlignVCenter,
+                        label_text,
+                    )
+                    painter.restore()
+
         painter.end()
 
         self._update_display()
@@ -787,7 +821,13 @@ class AnnotationCanvasWidget(QWidget):
             f"Drew saved polyline with {len(polyline)} points in color {color}"
         )
 
-    def draw_saved_bbox(self, bbox: List[float], color: str, width: int = 3):
+    def draw_saved_bbox(
+        self,
+        bbox: List[float],
+        color: str,
+        width: int = 3,
+        label: Optional[str] = None,
+    ):
         """
         Draw a bounding box from database coordinates onto the annotation canvas.
 
@@ -796,6 +836,7 @@ class AnnotationCanvasWidget(QWidget):
                   in normalized coordinates (0-1)
             color: Color hex string (e.g., '#FF0000')
             width: Line width in pixels
+            label: Optional text label to render near the bounding box
         """
         if not self.annotation_pixmap or not self.original_pixmap:
             logger.warning("Cannot draw bounding box: no image loaded")
@@ -828,11 +869,11 @@ class AnnotationCanvasWidget(QWidget):
         self.bboxes.append(
             [float(x_min_norm), float(y_min_norm), float(x_max_norm), float(y_max_norm)]
         )
-        self.bbox_meta.append({"color": pen_color, "width": int(width)})
+        self.bbox_meta.append({"color": pen_color, "width": int(width), "label": label})
 
         # Store in all_strokes for consistency
         self.all_strokes.append(
-            {"bbox": bbox, "color": color, "alpha": 128, "width": width}
+            {"bbox": bbox, "color": color, "alpha": 128, "width": width, "label": label}
         )
 
         # Redraw overlay (polylines + all bounding boxes)