saw_mill_knot_detection/tk_annotation_gui.py

#!/usr/bin/env python3
"""Tkinter-based annotation GUI.

This is a standalone GUI for manual bounding-box annotation that writes
`annotations.json` in the same format used by the project:

{
  "image.jpg": [
    {"bbox": [x1, y1, x2, y2], "label": "knot", "confidence": 1.0, "source": "manual"},
    ...
  ],
  ...
}

This project uses the Tkinter GUI as the annotation interface.

Run:
  python tk_annotation_gui.py

Optional:
  python tk_annotation_gui.py --images-dir IMAGE/

Controls:
- Click-drag on the image to create a box
- Double-click a box entry to delete it
- Prev/Next to navigate

Notes:
- Boxes are stored in ORIGINAL image pixel coordinates.
- The displayed image is scaled to fit the canvas; coordinates are converted.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

try:
    import torch
    CUDA_AVAILABLE = torch.cuda.is_available()
except ImportError:
    CUDA_AVAILABLE = False

import tkinter as tk
from tkinter import ttk

from PIL import Image, ImageTk

# Defaults
DEFAULT_IMAGES_DIR = "IMAGE/"
DEFAULT_MODEL_WEIGHTS = ""
ANNOTATION_CATEGORIES = ["knot"]
DEFAULT_DETECTION_THRESHOLD = 0.5

try:
    import config as _cfg

    DEFAULT_IMAGES_DIR = getattr(_cfg, "DEFAULT_IMAGES_DIR", DEFAULT_IMAGES_DIR)
    DEFAULT_MODEL_WEIGHTS = getattr(_cfg, "DEFAULT_MODEL_WEIGHTS", DEFAULT_MODEL_WEIGHTS)
    ANNOTATION_CATEGORIES = getattr(_cfg, "ANNOTATION_CATEGORIES", ANNOTATION_CATEGORIES)
    DEFAULT_DETECTION_THRESHOLD = float(getattr(_cfg, "DEFAULT_DETECTION_THRESHOLD", DEFAULT_DETECTION_THRESHOLD))
except Exception:
    pass


@dataclass
class DisplayTransform:
    scale: float
    offset_x: float
    offset_y: float


class TkAnnotationApp:
    def __init__(self, root: tk.Tk, images_dir: Path):
        self.root = root
        self.root.title("Wood Knot Annotation Tool (Tkinter)")

        self.images_dir = images_dir
        self.image_paths: list[Path] = []
        self.current_idx: int = 0

        self.ann_file: Path = self.images_dir / "annotations.json"
        self.annotations: dict[str, list[dict[str, Any]]] = {}

        self.current_image: Image.Image | None = None
        self.current_image_path: Path | None = None
        self.current_photo: ImageTk.PhotoImage | None = None

        self.transform: DisplayTransform | None = None

        self._draw_start: tuple[float, float] | None = None
        self._preview_rect_id: int | None = None

        # New variables for box selection and editing
        self.selected_box_index: int | None = None
        self.dragging: bool = False
        self.drag_start: tuple[float, float] | None = None
        self.drag_mode: str | None = None  # 'move' or 'resize'
        self.resize_corner: str | None = None  # 'nw', 'ne', 'sw', 'se'
        self._is_selecting: bool = False
        self._potential_select: int | None = None
        self._mouse_moved: bool = False

        self.model: Any | None = None
        self.model_type: str | None = None  # rf-detr | rt-detr | yolov6 | yolox
        self.model_path: Path | None = None

        self.model_path_var = tk.StringVar(value=str(DEFAULT_MODEL_WEIGHTS) if DEFAULT_MODEL_WEIGHTS else "")
        self.model_type_var = tk.StringVar(value="auto")
        self.threshold_var = tk.DoubleVar(value=float(DEFAULT_DETECTION_THRESHOLD))

        self.label_var = tk.StringVar(value=(ANNOTATION_CATEGORIES[0] if ANNOTATION_CATEGORIES else "knot"))

        self._build_ui()
        self._load_images_dir(self.images_dir)
        self._auto_load_model()

    def _auto_load_model(self) -> None:
        if DEFAULT_MODEL_WEIGHTS and Path(DEFAULT_MODEL_WEIGHTS).expanduser().exists():
            self._set_model_status("Auto-loading model...")
            self.load_model()

    # ------------------------- UI -------------------------

    def _build_ui(self) -> None:
        container = ttk.Frame(self.root, padding=8)
        container.grid(row=0, column=0, sticky="nsew")

        self.root.rowconfigure(0, weight=1)
        self.root.columnconfigure(0, weight=1)

        # Top controls
        top = ttk.Frame(container)
        top.grid(row=0, column=0, columnspan=2, sticky="ew", pady=(0, 8))
        top.columnconfigure(1, weight=1)

        ttk.Label(top, text="Images dir:").grid(row=0, column=0, sticky="w")
        self.images_dir_var = tk.StringVar(value=str(self.images_dir))
        self.images_dir_entry = ttk.Entry(top, textvariable=self.images_dir_var)
        self.images_dir_entry.grid(row=0, column=1, sticky="ew", padx=6)
        ttk.Button(top, text="Load", command=self._on_load_dir).grid(row=0, column=2, sticky="ew")

        self.index_label = ttk.Label(top, text="Image: -/-")
        self.index_label.grid(row=0, column=3, sticky="e", padx=(10, 0))

        ttk.Separator(container, orient="horizontal").grid(row=1, column=0, columnspan=2, sticky="ew", pady=(0, 8))

        # Left: Canvas
        left = ttk.Frame(container)
        left.grid(row=2, column=0, sticky="nsew", padx=(0, 8))
        container.rowconfigure(2, weight=1)
        container.columnconfigure(0, weight=3)

        nav = ttk.Frame(left)
        nav.grid(row=0, column=0, sticky="ew", pady=(0, 6))
        nav.columnconfigure(2, weight=1)

        ttk.Button(nav, text="Prev", command=self.prev_image).grid(row=0, column=0, sticky="w")
        ttk.Button(nav, text="Next", command=self.next_image).grid(row=0, column=1, sticky="w", padx=(6, 0))
        self.status_label = ttk.Label(nav, text="", foreground="#444")
        self.status_label.grid(row=0, column=2, sticky="w", padx=(10, 0))

        self.canvas = tk.Canvas(left, width=1200, height=800, bg="#111", highlightthickness=0)
        self.canvas.grid(row=1, column=0, sticky="nsew")
        left.rowconfigure(1, weight=1)
        left.columnconfigure(0, weight=1)

        self.canvas.bind("<ButtonPress-1>", self._on_mouse_down)
        self.canvas.bind("<B1-Motion>", self._on_mouse_move)
        self.canvas.bind("<ButtonRelease-1>", self._on_mouse_up)

        # New binds for right-click resize and delete key
        self.canvas.bind("<ButtonPress-3>", self._on_right_mouse_down)
        self.canvas.bind("<B3-Motion>", self._on_right_mouse_move)
        self.canvas.bind("<ButtonRelease-3>", self._on_right_mouse_up)
        self.root.bind("<Delete>", self._on_delete_key)

        # Right: boxes list + controls
        right = ttk.Frame(container)
        right.grid(row=2, column=1, sticky="nsew")
        container.columnconfigure(1, weight=1)

        # Model controls
        model_frame = ttk.LabelFrame(right, text="Auto-Label", padding=8)
        model_frame.grid(row=0, column=0, columnspan=2, sticky="ew")
        model_frame.columnconfigure(1, weight=1)

        ttk.Label(model_frame, text="Weights:").grid(row=0, column=0, sticky="w")
        self.model_entry = ttk.Entry(model_frame, textvariable=self.model_path_var)
        self.model_entry.grid(row=0, column=1, sticky="ew", padx=(6, 0))

        ttk.Label(model_frame, text="Type:").grid(row=1, column=0, sticky="w", pady=(6, 0))
        self.model_type_menu = ttk.OptionMenu(
            model_frame,
            self.model_type_var,
            self.model_type_var.get(),
            "auto",
            "rf-detr",
            "rt-detr",
            "yolov6",
            "yolox",
        )
        self.model_type_menu.grid(row=1, column=1, sticky="ew", padx=(6, 0), pady=(6, 0))

        ttk.Label(model_frame, text="Threshold:").grid(row=2, column=0, sticky="w", pady=(6, 0))
        self.threshold_scale = ttk.Scale(model_frame, from_=0.05, to=0.95, variable=self.threshold_var)
        self.threshold_scale.grid(row=2, column=1, sticky="ew", padx=(6, 0), pady=(6, 0))

        model_buttons = ttk.Frame(model_frame)
        model_buttons.grid(row=3, column=0, columnspan=2, sticky="ew", pady=(8, 0))
        model_buttons.columnconfigure(0, weight=1)
        model_buttons.columnconfigure(1, weight=1)
        ttk.Button(model_buttons, text="Load Model", command=self.load_model).grid(row=0, column=0, sticky="ew")
        ttk.Button(model_buttons, text="Auto-Label Current", command=self.auto_label_current).grid(row=0, column=1, sticky="ew", padx=(6, 0))

        self.model_status = ttk.Label(model_frame, text="No model loaded")
        self.model_status.grid(row=4, column=0, columnspan=2, sticky="w", pady=(6, 0))

        ttk.Label(right, text="Label:").grid(row=1, column=0, sticky="w", pady=(10, 0))
        self.label_menu = ttk.OptionMenu(right, self.label_var, self.label_var.get(), *ANNOTATION_CATEGORIES)
        self.label_menu.grid(row=1, column=1, sticky="ew", padx=(6, 0), pady=(10, 0))
        right.columnconfigure(1, weight=1)

        ttk.Label(right, text="Annotations:").grid(row=2, column=0, columnspan=2, sticky="w", pady=(10, 4))

        self.box_list = tk.Listbox(right, height=18)
        self.box_list.grid(row=3, column=0, columnspan=2, sticky="nsew")
        right.rowconfigure(3, weight=1)

        self.box_list.bind("<Double-Button-1>", self._on_box_double_click)
        self.box_list.bind("<<ListboxSelect>>", self._on_box_select)

        buttons = ttk.Frame(right)
        buttons.grid(row=4, column=0, columnspan=2, sticky="ew", pady=(6, 0))
        ttk.Button(buttons, text="Delete Selected", command=self.delete_selected_box).grid(row=0, column=0, sticky="ew")
        ttk.Button(buttons, text="Clear All", command=self.clear_all_boxes).grid(row=0, column=1, sticky="ew", padx=(6, 0))

        # Make buttons frame expand reasonably
        buttons.columnconfigure(0, weight=1)
        buttons.columnconfigure(1, weight=1)

    # ------------------------- Model loading / auto-label -------------------------

    def _guess_model_type_from_path(self, path: Path) -> str:
        s = str(path).lower()
        if "rf" in s or "checkpoint" in s or s.endswith(".pth"):
            return "rf-detr"
        if "rtdetr" in s or "rt-detr" in s:
            return "rt-detr"
        if "yolov6" in s:
            return "yolov6"
        if "yolox" in s:
            return "yolox"
        # Default to ultralytics RT-DETR if ambiguous
        return "rt-detr"

    def load_model(self) -> None:
        raw = self.model_path_var.get().strip()
        if not raw:
            self._set_model_status("No weights path provided")
            return

        weights_path = Path(raw).expanduser()
        if not weights_path.exists():
            self._set_model_status(f"File not found: {weights_path}")
            return

        selected = (self.model_type_var.get() or "auto").strip().lower()
        model_type = self._guess_model_type_from_path(weights_path) if selected == "auto" else selected

        try:
            # RF-DETR optional
            if model_type == "rf-detr":
                from rfdetr import RFDETRNano

                self.model = RFDETRNano(pretrain_weights=str(weights_path))
            else:
                if model_type == "rt-detr":
                    from ultralytics import RTDETR

                    self.model = RTDETR(str(weights_path))
                else:
                    from ultralytics import YOLO

                    self.model = YOLO(str(weights_path))

            self.model_type = model_type
            self.model_path = weights_path
            self._set_model_status(f"Loaded: {weights_path.name} ({model_type})")
        except Exception as e:
            self.model = None
            self.model_type = None
            self.model_path = None
            self._set_model_status(f"Load failed: {e}")

    def auto_label_current(self) -> None:
        if self.current_image_path is None:
            return
        if self.model is None or self.model_type is None:
            self._set_model_status("No model loaded")
            return

        threshold = float(self.threshold_var.get())
        img_path = self.current_image_path

        try:
            new_boxes: list[dict[str, Any]] = []

            if self.model_type == "rf-detr":
                # RF-DETR model expects PIL image
                if self.current_image is None:
                    return
                detections = self.model.predict(self.current_image, threshold=threshold)
                for i in range(len(detections)):
                    xyxy = detections.xyxy[i]
                    conf = float(detections.confidence[i]) if detections.confidence is not None else 1.0
                    x1, y1, x2, y2 = xyxy
                    new_boxes.append(
                        {
                            "bbox": [float(x1), float(y1), float(x2), float(y2)],
                            "label": "knot",
                            "confidence": conf,
                            "source": "auto",
                        }
                    )
            else:
                # Ultralytics models
                device = 'cuda' if CUDA_AVAILABLE else 'cpu'
                results = self.model.predict(source=self.current_image, conf=threshold, save=False, verbose=False, device=device)
                for result in results:
                    for box in result.boxes:
                        x1, y1, x2, y2 = box.xyxy[0].tolist()
                        conf = float(box.conf[0])
                        label = "knot"
                        try:
                            cls = int(box.cls[0])
                            if hasattr(self.model, "names") and cls in self.model.names:
                                label = str(self.model.names[cls])
                        except Exception:
                            pass

                        new_boxes.append(
                            {
                                "bbox": [float(x1), float(y1), float(x2), float(y2)],
                                "label": label,
                                "confidence": conf,
                                "source": "auto",
                            }
                        )

            # Match legacy behavior: append auto boxes to existing
            key = img_path.name
            # Remove previous auto labels
            existing_boxes = self.annotations.get(key, [])
            self.annotations[key] = [box for box in existing_boxes if box.get("source") != "auto"]
            self.annotations[key].extend(new_boxes)
            self._save_annotations()

            self._refresh_box_list()
            self._redraw_boxes()
            self._set_model_status(f"Auto-labeled: {len(new_boxes)}")
        except Exception as e:
            self._set_model_status(f"Auto-label failed: {e}")

    def _set_model_status(self, msg: str) -> None:
        self.model_status.config(text=msg)

    # ------------------------- Data load/save -------------------------

    def _load_images_dir(self, images_dir: Path) -> None:
        images_dir = images_dir.expanduser().resolve()
        if not images_dir.exists() or not images_dir.is_dir():
            self._set_status(f"Invalid images dir: {images_dir}")
            return

        self.images_dir = images_dir
        self.ann_file = self.images_dir / "annotations.json"

        self.image_paths = sorted(list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.png")) + list(images_dir.glob("*.jpeg")))
        self.current_idx = 0

        # Load annotations (if present)
        self.annotations = {}
        if self.ann_file.exists():
            try:
                with self.ann_file.open("r") as f:
                    data = json.load(f)
                if isinstance(data, dict):
                    self.annotations = data
            except Exception as e:
                self._set_status(f"Failed to load annotations.json: {e}")

        if not self.image_paths:
            self._set_status("No images found")
            self._clear_canvas()
            self._update_index_label()
            self._refresh_box_list()
            return

        self._set_status("")
        self.load_current_image()

    def _save_annotations(self) -> None:
        # Ensure we always have an entry for current image
        if self.current_image_path is not None:
            key = self.current_image_path.name
            self.annotations.setdefault(key, [])

        try:
            with self.ann_file.open("w") as f:
                json.dump(self.annotations, f, indent=2)
        except Exception as e:
            self._set_status(f"Failed to save annotations: {e}")

    # ------------------------- Navigation -------------------------

    def prev_image(self) -> None:
        if not self.image_paths:
            return
        self.current_idx = max(0, self.current_idx - 1)
        self.load_current_image()

    def next_image(self) -> None:
        if not self.image_paths:
            return
        self.current_idx = min(len(self.image_paths) - 1, self.current_idx + 1)
        self.load_current_image()

    def load_current_image(self) -> None:
        if not self.image_paths:
            return

        self.current_image_path = self.image_paths[self.current_idx]
        self.selected_box_index = None  # Reset selection

        try:
            img = Image.open(self.current_image_path).convert("RGB")
        except Exception as e:
            self._set_status(f"Failed to open image: {e}")
            return

        self.current_image = img
        self._update_index_label()

        # Ensure annotation list exists
        self.annotations.setdefault(self.current_image_path.name, [])

        self._render_image_and_boxes()
        self._refresh_box_list()

    def _update_index_label(self) -> None:
        total = len(self.image_paths)
        if total == 0:
            self.index_label.config(text="Image: -/-")
            return
        filename = self.image_paths[self.current_idx].name
        self.index_label.config(text=f"Image {self.current_idx + 1}/{total}: {filename}")

    # ------------------------- Canvas rendering -------------------------

    def _clear_canvas(self) -> None:
        self.canvas.delete("all")
        self.current_photo = None
        self.transform = None

    def _render_image_and_boxes(self) -> None:
        self._clear_canvas()
        if self.current_image is None:
            return

        canvas_w = int(self.canvas.winfo_width())
        canvas_h = int(self.canvas.winfo_height())
        # If not yet realized, fall back to configured size
        if canvas_w <= 2:
            canvas_w = int(self.canvas["width"])
        if canvas_h <= 2:
            canvas_h = int(self.canvas["height"])

        orig_w, orig_h = self.current_image.size
        scale = min(canvas_w / orig_w, canvas_h / orig_h)
        scale = max(scale, 1e-6)

        disp_w = int(orig_w * scale)
        disp_h = int(orig_h * scale)

        offset_x = (canvas_w - disp_w) / 2
        offset_y = (canvas_h - disp_h) / 2

        disp_img = self.current_image.resize((disp_w, disp_h), Image.Resampling.BILINEAR)
        self.current_photo = ImageTk.PhotoImage(disp_img)

        # Draw image
        self.canvas.create_image(offset_x, offset_y, anchor="nw", image=self.current_photo)
        self.transform = DisplayTransform(scale=scale, offset_x=offset_x, offset_y=offset_y)

        # Draw boxes
        self._redraw_boxes()

    def _redraw_boxes(self) -> None:
        self.canvas.delete("box")
        if self.current_image_path is None or self.transform is None:
            return

        boxes = self.annotations.get(self.current_image_path.name, []) or []
        for i, box in enumerate(boxes):
            bbox = box.get("bbox") if isinstance(box, dict) else None
            if not bbox or len(bbox) != 4:
                continue
            x1, y1, x2, y2 = bbox
            dx1, dy1 = self._img_to_disp(x1, y1)
            dx2, dy2 = self._img_to_disp(x2, y2)
            color = "#FF4444" if i == self.selected_box_index else "#00FF66"
            width = 3 if i == self.selected_box_index else 2
            self.canvas.create_rectangle(dx1, dy1, dx2, dy2, outline=color, width=width, tags=("box", f"box_{i}"))
            if i == self.selected_box_index:
                # Draw resize handles
                handle_size = 6
                self.canvas.create_rectangle(dx1-handle_size, dy1-handle_size, dx1+handle_size, dy1+handle_size, fill=color, tags=("box", f"box_{i}"))
                self.canvas.create_rectangle(dx2-handle_size, dy1-handle_size, dx2+handle_size, dy1+handle_size, fill=color, tags=("box", f"box_{i}"))
                self.canvas.create_rectangle(dx1-handle_size, dy2-handle_size, dx1+handle_size, dy2+handle_size, fill=color, tags=("box", f"box_{i}"))
                self.canvas.create_rectangle(dx2-handle_size, dy2-handle_size, dx2+handle_size, dy2+handle_size, fill=color, tags=("box", f"box_{i}"))

    def _img_to_disp(self, x: float, y: float) -> tuple[float, float]:
        assert self.transform is not None
        return (x * self.transform.scale + self.transform.offset_x, y * self.transform.scale + self.transform.offset_y)

    def _disp_to_img(self, x: float, y: float) -> tuple[float, float]:
        assert self.transform is not None
        ix = (x - self.transform.offset_x) / self.transform.scale
        iy = (y - self.transform.offset_y) / self.transform.scale
        if self.current_image is None:
            return ix, iy
        w, h = self.current_image.size
        ix = min(max(ix, 0.0), float(w))
        iy = min(max(iy, 0.0), float(h))
        return ix, iy

    def _find_box_at_point(self, x: float, y: float) -> int | None:
        """Find the box at the given display coordinates, prioritizing smaller boxes."""
        if self.current_image_path is None:
            return None
        boxes = self.annotations.get(self.current_image_path.name, []) or []
        candidates = []
        for i, box in enumerate(boxes):
            bbox = box.get("bbox")
            if not bbox or len(bbox) != 4:
                continue
            x1, y1, x2, y2 = bbox
            dx1, dy1 = self._img_to_disp(x1, y1)
            dx2, dy2 = self._img_to_disp(x2, y2)
            if dx1 <= x <= dx2 and dy1 <= y <= dy2:
                area = (x2 - x1) * (y2 - y1)
                candidates.append((area, i))
        if not candidates:
            return None
        # Sort by area ascending (smaller first)
        candidates.sort()
        return candidates[0][1]

    def _find_resize_corner(self, x: float, y: float, box_index: int) -> str | None:
        """Find which corner/handle is clicked for resizing."""
        if self.current_image_path is None:
            return None
        boxes = self.annotations.get(self.current_image_path.name, []) or []
        if box_index >= len(boxes):
            return None
        bbox = boxes[box_index].get("bbox")
        if not bbox or len(bbox) != 4:
            return None
        x1, y1, x2, y2 = bbox
        dx1, dy1 = self._img_to_disp(x1, y1)
        dx2, dy2 = self._img_to_disp(x2, y2)
        handle_size = 10  # Slightly larger for easier clicking
        corners = {
            'nw': (dx1, dy1),
            'ne': (dx2, dy1),
            'sw': (dx1, dy2),
            'se': (dx2, dy2)
        }
        for corner, (cx, cy) in corners.items():
            if cx - handle_size <= x <= cx + handle_size and cy - handle_size <= y <= cy + handle_size:
                return corner
        return None

    def _find_handle_at_point(self, x: float, y: float) -> tuple[int, str] | None:
        """Return (box_index, corner) if point hits a resize handle.

        Unlike _find_box_at_point, this can hit even slightly outside the box
        bounds (since the drawn handle squares extend beyond the rectangle).
        When handles overlap, choose the smallest box by area.
        """

        if self.current_image_path is None:
            return None
        boxes = self.annotations.get(self.current_image_path.name, []) or []

        handle_size = 10
        candidates: list[tuple[float, int, str]] = []  # (area, index, corner)
        for i, box in enumerate(boxes):
            bbox = box.get("bbox") if isinstance(box, dict) else None
            if not bbox or len(bbox) != 4:
                continue
            x1, y1, x2, y2 = bbox
            area = (x2 - x1) * (y2 - y1)
            dx1, dy1 = self._img_to_disp(x1, y1)
            dx2, dy2 = self._img_to_disp(x2, y2)
            corners = {
                "nw": (dx1, dy1),
                "ne": (dx2, dy1),
                "sw": (dx1, dy2),
                "se": (dx2, dy2),
            }
            for corner, (cx, cy) in corners.items():
                if cx - handle_size <= x <= cx + handle_size and cy - handle_size <= y <= cy + handle_size:
                    candidates.append((area, i, corner))

        if not candidates:
            return None
        candidates.sort(key=lambda t: t[0])
        _, idx, corner = candidates[0]
        return idx, corner

    def _apply_resize_drag(self, event_x: float, event_y: float) -> None:
        if (
            not self.dragging
            or self.drag_mode != "resize"
            or self.resize_corner is None
            or self.selected_box_index is None
            or self.drag_start is None
            or self.current_image_path is None
        ):
            return

        boxes = self.annotations.get(self.current_image_path.name, []) or []
        if self.selected_box_index >= len(boxes):
            return

        bbox = boxes[self.selected_box_index].get("bbox")
        if not bbox or len(bbox) != 4:
            return

        x1, y1, x2, y2 = bbox
        dx = event_x - self.drag_start[0]
        dy = event_y - self.drag_start[1]

        # Convert to display coords
        dx1, dy1 = self._img_to_disp(x1, y1)
        dx2, dy2 = self._img_to_disp(x2, y2)

        if "n" in self.resize_corner:
            dy1 += dy
        if "s" in self.resize_corner:
            dy2 += dy
        if "w" in self.resize_corner:
            dx1 += dx
        if "e" in self.resize_corner:
            dx2 += dx

        # Convert back to image coords
        ix1, iy1 = self._disp_to_img(dx1, dy1)
        ix2, iy2 = self._disp_to_img(dx2, dy2)

        # Ensure min size
        if abs(ix2 - ix1) < 2:
            ix2 = ix1 + 2 if ix2 > ix1 else ix1 - 2
        if abs(iy2 - iy1) < 2:
            iy2 = iy1 + 2 if iy2 > iy1 else iy1 - 2

        boxes[self.selected_box_index]["bbox"] = [min(ix1, ix2), min(iy1, iy2), max(ix1, ix2), max(iy1, iy2)]
        self.drag_start = (event_x, event_y)
        self._redraw_boxes()

    # ------------------------- Mouse interactions -------------------------

    def _on_mouse_down(self, event: tk.Event) -> None:
        if self.current_image is None or self.current_image_path is None or self.transform is None:
            return

        # First priority: if cursor is over a corner handle, start resizing (left-drag)
        handle_hit = self._find_handle_at_point(event.x, event.y)
        if handle_hit is not None:
            box_index, corner = handle_hit
            self.selected_box_index = box_index
            self.dragging = True
            self.drag_mode = "resize"
            self.resize_corner = corner
            self.drag_start = (event.x, event.y)
            self._is_selecting = False
            self._potential_select = None
            self._mouse_moved = False
            self._refresh_box_list()
            self._redraw_boxes()
            return
        
        # Find which box (if any) is under the cursor
        box_index = self._find_box_at_point(event.x, event.y)
        
        # Check if Ctrl is held for moving or resizing boxes
        if event.state & 0x4:  # Ctrl key
            # First, check if clicking on a corner of the selected box for resizing
            if self.selected_box_index is not None:
                corner = self._find_resize_corner(event.x, event.y, self.selected_box_index)
                if corner:
                    self.dragging = True
                    self.drag_mode = 'resize'
                    self.resize_corner = corner
                    self.drag_start = (event.x, event.y)
                    self._refresh_box_list()
                    self._redraw_boxes()
                    return
            
            # Otherwise, select and move a box
            if box_index is not None:
                self.selected_box_index = box_index
                self.dragging = True
                self.drag_mode = 'move'
                self.drag_start = (event.x, event.y)
                self._refresh_box_list()
                self._redraw_boxes()
                return

        # Normal mode: check if clicking inside a box to potentially select it
        if box_index is not None:
            self._potential_select = box_index
            self._is_selecting = True
            self._mouse_moved = False
            return
        
        # Otherwise, start drawing
        self._draw_start = (event.x, event.y)
        self._is_selecting = False
        if self._preview_rect_id is not None:
            self.canvas.delete(self._preview_rect_id)
            self._preview_rect_id = None

    def _on_mouse_move(self, event: tk.Event) -> None:
        if self.dragging and self.drag_mode == "resize":
            self._apply_resize_drag(event.x, event.y)
            return

        if self.dragging and self.drag_mode == 'move' and self.drag_start and self.selected_box_index is not None:
            # Move the box
            dx = event.x - self.drag_start[0]
            dy = event.y - self.drag_start[1]
            if self.current_image_path is None:
                return
            boxes = self.annotations.get(self.current_image_path.name, []) or []
            if self.selected_box_index >= len(boxes):
                return
            bbox = boxes[self.selected_box_index]["bbox"]
            x1, y1, x2, y2 = bbox
            # Convert to display coords, move, convert back
            dx1, dy1 = self._img_to_disp(x1, y1)
            dx2, dy2 = self._img_to_disp(x2, y2)
            dx1 += dx
            dy1 += dy
            dx2 += dx
            dy2 += dy
            ix1, iy1 = self._disp_to_img(dx1, dy1)
            ix2, iy2 = self._disp_to_img(dx2, dy2)
            boxes[self.selected_box_index]["bbox"] = [ix1, iy1, ix2, iy2]
            self.drag_start = (event.x, event.y)
            self._redraw_boxes()
            return
        
        if self._is_selecting:
            self._mouse_moved = True
            return
        
        if self._draw_start is None or self.current_image is None or self.transform is None:
            return

        x0, y0 = self._draw_start
        x1, y1 = event.x, event.y

        if self._preview_rect_id is not None:
            self.canvas.delete(self._preview_rect_id)

        self._preview_rect_id = self.canvas.create_rectangle(
            x0, y0, x1, y1, outline="#FFCC00", width=2, dash=(4, 2)
        )

    def _on_mouse_up(self, event: tk.Event) -> None:
        if self.dragging:
            self.dragging = False
            self.drag_mode = None
            self.drag_start = None
            self.resize_corner = None
            self._save_annotations()
            self._refresh_box_list()
            self._redraw_boxes()
            return
        
        if self._is_selecting:
            if not self._mouse_moved and self._potential_select is not None:
                self.selected_box_index = self._potential_select
                self._refresh_box_list()
                self._redraw_boxes()
            self._is_selecting = False
            self._potential_select = None
            self._mouse_moved = False
            return
        
        if self._draw_start is None or self.current_image is None or self.current_image_path is None or self.transform is None:
            self._draw_start = None
            return

        x0, y0 = self._draw_start
        x1, y1 = event.x, event.y
        self._draw_start = None

        if self._preview_rect_id is not None:
            self.canvas.delete(self._preview_rect_id)
            self._preview_rect_id = None

        # Convert to image coords
        ix0, iy0 = self._disp_to_img(x0, y0)
        ix1, iy1 = self._disp_to_img(x1, y1)

        x_min, x_max = sorted([ix0, ix1])
        y_min, y_max = sorted([iy0, iy1])

        # Ignore tiny drags
        if (x_max - x_min) < 2 or (y_max - y_min) < 2:
            return

        new_box = {
            "bbox": [float(x_min), float(y_min), float(x_max), float(y_max)],
            "label": self.label_var.get() or "knot",
            "confidence": 1.0,
            "source": "manual",
        }

        boxes = self.annotations.setdefault(self.current_image_path.name, [])
        boxes.append(new_box)
        self._save_annotations()

        self._refresh_box_list()
        self._redraw_boxes()

    def _on_right_mouse_down(self, event: tk.Event) -> None:
        if self.current_image is None or self.current_image_path is None or self.transform is None:
            return
        box_index = self._find_box_at_point(event.x, event.y)
        if box_index is not None:
            corner = self._find_resize_corner(event.x, event.y, box_index)
            if corner:
                self.selected_box_index = box_index
                self.dragging = True
                self.drag_mode = 'resize'
                self.resize_corner = corner
                self.drag_start = (event.x, event.y)
                self._refresh_box_list()
                self._redraw_boxes()

    def _on_right_mouse_move(self, event: tk.Event) -> None:
        self._apply_resize_drag(event.x, event.y)

    def _on_right_mouse_up(self, event: tk.Event) -> None:
        if self.dragging and self.drag_mode == 'resize':
            self.dragging = False
            self.drag_mode = None
            self.resize_corner = None
            self.drag_start = None
            self._save_annotations()
            self._refresh_box_list()
            self._redraw_boxes()

    def _on_delete_key(self, event: tk.Event) -> None:
        self.delete_selected_box()

    def _refresh_box_list(self) -> None:
        self.box_list.delete(0, tk.END)
        if self.current_image_path is None:
            return
        boxes = self.annotations.get(self.current_image_path.name, []) or []
        for idx, box in enumerate(boxes):
            if not isinstance(box, dict) or "bbox" not in box:
                continue
            x1, y1, x2, y2 = box["bbox"]
            label = str(box.get("label", "knot"))
            src = str(box.get("source", "manual"))
            conf = box.get("confidence", 1.0)
            marker = "[x]" if idx == self.selected_box_index else "[ ]"
            self.box_list.insert(
                tk.END,
                f"{marker} {idx}: {label} ({src}, {conf:.3f})  ({x1:.1f},{y1:.1f})-({x2:.1f},{y2:.1f})",
            )
        # Select the item in listbox if selected
        if self.selected_box_index is not None and self.selected_box_index < self.box_list.size():
            self.box_list.selection_set(self.selected_box_index)

    def _selected_box_index(self) -> int | None:
        sel = self.box_list.curselection()
        if not sel:
            return None
        idx = int(sel[0])
        self.selected_box_index = idx
        self._redraw_boxes()
        return idx

    def delete_selected_box(self) -> None:
        if self.current_image_path is None:
            return
        idx = self._selected_box_index()
        if idx is None:
            return

        boxes = self.annotations.get(self.current_image_path.name, []) or []
        if 0 <= idx < len(boxes):
            del boxes[idx]
            if self.selected_box_index == idx:
                self.selected_box_index = None
            elif self.selected_box_index is not None and self.selected_box_index > idx:
                self.selected_box_index -= 1
            self._save_annotations()
            self._refresh_box_list()
            self._redraw_boxes()
    def _on_box_select(self, event: tk.Event) -> None:
        self._selected_box_index()
        self._refresh_box_list()  # To update markers
    def _on_box_double_click(self, _event: tk.Event) -> None:
        self.delete_selected_box()

    def clear_all_boxes(self) -> None:
        if self.current_image_path is None:
            return
        self.annotations[self.current_image_path.name] = []
        self.selected_box_index = None
        self._save_annotations()
        self._refresh_box_list()
        self._redraw_boxes()

    # ------------------------- Misc -------------------------

    def _set_status(self, msg: str) -> None:
        self.status_label.config(text=msg)

    def _on_load_dir(self) -> None:
        path = Path(self.images_dir_var.get().strip())
        self._load_images_dir(path)


def main() -> None:
    parser = argparse.ArgumentParser(description="Tkinter annotation GUI")
    parser.add_argument(
        "--images-dir",
        type=Path,
        default=Path(DEFAULT_IMAGES_DIR) if DEFAULT_IMAGES_DIR else Path("IMAGE/"),
        help="Directory containing images and annotations.json",
    )
    args = parser.parse_args()

    root = tk.Tk()
    app = TkAnnotationApp(root, args.images_dir)

    # Re-render on first layout so scaling is correct
    def after_layout() -> None:
        if app.current_image is not None:
            app._render_image_and_boxes()

    root.after(50, after_layout)
    root.mainloop()


if __name__ == "__main__":
    main()