saw_mill_knot_detection/annotation_gui.py

"""
Simple customizable annotation GUI with auto-labeling support.

Built with Gradio - easy to modify and extend.
Run: python annotation_gui.py

To set default paths, edit config.py
"""

from __future__ import annotations

import argparse
import html
import json
import subprocess
import threading
from pathlib import Path
from typing import Any

import gradio as gr
from PIL import Image, ImageDraw

# Try to load config, use fallbacks if not available
try:
    from config import (
        DEFAULT_IMAGES_DIR, DEFAULT_MODEL_WEIGHTS, DEFAULT_PORT,
        DEFAULT_DETECTION_THRESHOLD, DEFAULT_TRAIN_EPOCHS,
        DEFAULT_BATCH_SIZE, DEFAULT_LEARNING_RATE, DEFAULT_MODEL_SIZE
    )
except ImportError:
    DEFAULT_IMAGES_DIR = None
    DEFAULT_MODEL_WEIGHTS = None
    DEFAULT_PORT = 7860
    DEFAULT_DETECTION_THRESHOLD = 0.5
    DEFAULT_TRAIN_EPOCHS = 20
    DEFAULT_BATCH_SIZE = 4
    DEFAULT_LEARNING_RATE = 1e-4
    DEFAULT_MODEL_SIZE = "small"


# Gradio 6 sanitizes <script> tags inside gr.HTML content. Also, js_on_load only
# runs when the component is first mounted, not when its HTML updates. We
# therefore install a global initializer (via demo.launch(js=...)) and have
# js_on_load call into it when available.
CANVAS_JS_ON_LOAD = r"""
(() => {
    if (window.__initAnnotationCanvas) {
        window.__initAnnotationCanvas(element);
    }
})();
"""

CANVAS_GLOBAL_JS = r"""
(() => {
    function init(root) {
        if (!root) return;
        const canvas = root.querySelector('#annotation-canvas');
        const imgEl = root.querySelector('#annotation-img');
        const initialBoxesEl = root.querySelector('#annotation-initial-boxes');
        if (!canvas || !imgEl || !initialBoxesEl) return;

        // Ensure we don't double-bind if Gradio reuses the DOM node.
        if (canvas.dataset.bound === '1') {
            // Still redraw in case boxes were updated.
            if (canvas.__redraw) canvas.__redraw();
            return;
        }
        canvas.dataset.bound = '1';

        const ctx = canvas.getContext('2d');
        const displayWidth = canvas.width;
        const displayHeight = canvas.height;

        let boxes = [];
        try {
            const raw = initialBoxesEl.value || initialBoxesEl.textContent || '[]';
            boxes = JSON.parse(raw);
            if (!Array.isArray(boxes)) boxes = [];
        } catch (_) {
            boxes = [];
        }

        const hiddenRoot = document.getElementById('canvas-boxes-data');
        const hiddenInput = hiddenRoot
            ? (hiddenRoot.querySelector('textarea, input') || hiddenRoot)
            : null;
        const syncHidden = () => {
            if (!hiddenInput) return;
            if (!('value' in hiddenInput)) return;
            hiddenInput.value = JSON.stringify(boxes);
            hiddenInput.dispatchEvent(new Event('input', { bubbles: true }));
        };
        syncHidden();

        let isDragging = false;
        let dragStart = null;
        let selectedCorner = null;
        let selectedBoxIndex = -1;
        let creatingBox = false;
        let createStart = null;

        function redraw() {
            ctx.clearRect(0, 0, displayWidth, displayHeight);
            // Base image is rendered via <img> below the canvas.

            boxes.forEach((box) => {
                const [x1, y1, x2, y2] = box.bbox;
                const label = box.label || 'knot';
                const conf = box.confidence || 1.0;

                ctx.strokeStyle = 'red';
                ctx.lineWidth = 3;
                ctx.strokeRect(x1, y1, x2 - x1, y2 - y1);

                const handleSize = 6;
                ctx.fillStyle = 'red';
                ctx.fillRect(x1 - handleSize, y1 - handleSize, handleSize * 2, handleSize * 2);
                ctx.fillRect(x2 - handleSize, y1 - handleSize, handleSize * 2, handleSize * 2);
                ctx.fillRect(x1 - handleSize, y2 - handleSize, handleSize * 2, handleSize * 2);
                ctx.fillRect(x2 - handleSize, y2 - handleSize, handleSize * 2, handleSize * 2);

                ctx.fillStyle = 'red';
                ctx.font = '16px Arial';
                const text = conf < 1.0 ? `${label} ${conf.toFixed(2)}` : label;
                ctx.fillText(text, x1, y1 - 5);
            });

            if (creatingBox && createStart && dragStart) {
                ctx.strokeStyle = 'blue';
                ctx.lineWidth = 2;
                ctx.setLineDash([5, 5]);
                const x = Math.min(createStart.x, dragStart.x);
                const y = Math.min(createStart.y, dragStart.y);
                const w = Math.abs(createStart.x - dragStart.x);
                const h = Math.abs(createStart.y - dragStart.y);
                ctx.strokeRect(x, y, w, h);
                ctx.setLineDash([]);
            }
        }
        canvas.__redraw = redraw;

        function getCornerAt(x, y) {
            const handleSize = 6;
            for (let i = 0; i < boxes.length; i++) {
                const [x1, y1, x2, y2] = boxes[i].bbox;
                const corners = [
                    { x: x1, y: y1, type: 'top-left' },
                    { x: x2, y: y1, type: 'top-right' },
                    { x: x1, y: y2, type: 'bottom-left' },
                    { x: x2, y: y2, type: 'bottom-right' },
                ];
                for (const corner of corners) {
                    if (
                        x >= corner.x - handleSize &&
                        x <= corner.x + handleSize &&
                        y >= corner.y - handleSize &&
                        y <= corner.y + handleSize
                    ) {
                        return { boxIndex: i, corner: corner.type, pos: corner };
                    }
                }
            }
            return null;
        }

        imgEl.addEventListener(
            'error',
            () => {
                ctx.clearRect(0, 0, displayWidth, displayHeight);
                ctx.fillStyle = '#ffcccc';
                ctx.fillRect(0, 0, displayWidth, displayHeight);
                ctx.fillStyle = 'black';
                ctx.font = '16px Arial';
                ctx.fillText('Image failed to load', 10, 30);
            },
            { once: true }
        );

        redraw();

        canvas.addEventListener('mousedown', (e) => {
            const rect = canvas.getBoundingClientRect();
            const x = (e.clientX - rect.left) * (displayWidth / rect.width);
            const y = (e.clientY - rect.top) * (displayHeight / rect.height);

            selectedCorner = getCornerAt(x, y);
            if (selectedCorner) {
                isDragging = true;
                selectedBoxIndex = selectedCorner.boxIndex;
                canvas.style.cursor = 'move';
            } else {
                creatingBox = true;
                createStart = { x, y };
                dragStart = { x, y };
                canvas.style.cursor = 'crosshair';
            }
        });

        canvas.addEventListener('mousemove', (e) => {
            const rect = canvas.getBoundingClientRect();
            const x = (e.clientX - rect.left) * (displayWidth / rect.width);
            const y = (e.clientY - rect.top) * (displayHeight / rect.height);

            if (isDragging && selectedCorner) {
                const box = boxes[selectedBoxIndex];
                if (selectedCorner.corner === 'top-left') {
                    box.bbox[0] = Math.min(x, box.bbox[2] - 10);
                    box.bbox[1] = Math.min(y, box.bbox[3] - 10);
                } else if (selectedCorner.corner === 'top-right') {
                    box.bbox[2] = Math.max(x, box.bbox[0] + 10);
                    box.bbox[1] = Math.min(y, box.bbox[3] - 10);
                } else if (selectedCorner.corner === 'bottom-left') {
                    box.bbox[0] = Math.min(x, box.bbox[2] - 10);
                    box.bbox[3] = Math.max(y, box.bbox[1] + 10);
                } else if (selectedCorner.corner === 'bottom-right') {
                    box.bbox[2] = Math.max(x, box.bbox[0] + 10);
                    box.bbox[3] = Math.max(y, box.bbox[1] + 10);
                }
                syncHidden();
                redraw();
                return;
            }

            if (creatingBox && createStart) {
                dragStart = { x, y };
                redraw();
                return;
            }

            const corner = getCornerAt(x, y);
            canvas.style.cursor = corner ? 'move' : 'crosshair';
        });

        canvas.addEventListener('mouseup', () => {
            if (creatingBox && createStart && dragStart) {
                const x1 = Math.min(createStart.x, dragStart.x);
                const y1 = Math.min(createStart.y, dragStart.y);
                const x2 = Math.max(createStart.x, dragStart.x);
                const y2 = Math.max(createStart.y, dragStart.y);
                const w = x2 - x1;
                const h = y2 - y1;
                if (w > 10 && h > 10) {
                    boxes.push({ bbox: [x1, y1, x2, y2], label: 'knot', confidence: 1.0, source: 'manual' });
                } else {
                    // Click without drag: create a default-size box around the click.
                    const size = 120;
                    const cx = createStart.x;
                    const cy = createStart.y;
                    const bx1 = Math.max(0, cx - size / 2);
                    const by1 = Math.max(0, cy - size / 2);
                    const bx2 = Math.min(displayWidth, cx + size / 2);
                    const by2 = Math.min(displayHeight, cy + size / 2);
                    boxes.push({ bbox: [bx1, by1, bx2, by2], label: 'knot', confidence: 1.0, source: 'manual' });
                }
                syncHidden();
                redraw();
            }

            isDragging = false;
            creatingBox = false;
            selectedCorner = null;
            selectedBoxIndex = -1;
            createStart = null;
            dragStart = null;
            canvas.style.cursor = 'crosshair';
        });
    }

    window.__initAnnotationCanvas = init;

    function scan() {
        document.querySelectorAll('[data-annotation-root="1"]').forEach((root) => init(root));
    }

    const obs = new MutationObserver(() => scan());
    obs.observe(document.documentElement, { childList: true, subtree: true });
    window.addEventListener('load', () => scan());
  scan();
  setTimeout(() => scan(), 0);
  setTimeout(() => scan(), 100);
  setTimeout(() => scan(), 250);
  setTimeout(() => scan(), 500);
  setTimeout(() => scan(), 1000);
  setTimeout(() => scan(), 2000);
})();
"""


class AnnotationApp:
    def __init__(self, images_dir: Path | None = None, model_weights: Path | None = None):
        self.images_dir = images_dir if images_dir else Path.cwd()
        self.current_model_path = model_weights
        self.current_model_type = None  # Track model type: 'rf-detr', 'rt-detr', 'yolov6', 'yolox'
        self.available_models = []  # List of discovered models for quick switching
        self.image_paths = []
        self.current_idx = 0
        self.annotations = {}  # image_name -> list of boxes
        self.model = None
        
        # Load images if directory provided
        if images_dir and images_dir.exists():
            self._load_images(images_dir)
        
        if model_weights and model_weights.exists():
            self._load_model(model_weights)
    
    def _load_images(self, images_dir: Path):
        """Load images from directory."""
        self.images_dir = images_dir
        self.image_paths = sorted(
            list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.png"))
        )
        self.current_idx = 0
        
        # Load existing annotations if present
        self.ann_file = images_dir / "annotations.json"
        if self.ann_file.exists():
            with self.ann_file.open("r") as f:
                self.annotations = json.load(f)
        else:
            self.annotations = {}
        
        return f"✓ Loaded {len(self.image_paths)} images from {images_dir}"
    
    def find_best_weights(self, directory: Path) -> tuple[Path | None, str | None]:
        """Find the best weights file in a directory based on model type detection."""
        if not directory.exists():
            return None, None
            
        # Check for RF-DETR weights (checkpoint_best_total.pth)
        rf_detr_weights = directory / "checkpoint_best_total.pth"
        if rf_detr_weights.exists():
            return rf_detr_weights, "rf-detr"
            
        # Check for Ultralytics weights (best.pt) in weights/ subdirectory
        ultralytics_weights = directory / "weights" / "best.pt"
        if ultralytics_weights.exists():
            # Try to determine specific type from directory name or other clues
            dir_name = directory.name.lower()
            if "rtdetr" in dir_name:
                return ultralytics_weights, "rt-detr"
            elif "yolov6" in dir_name:
                return ultralytics_weights, "yolov6"
            elif "yolox" in dir_name:
                return ultralytics_weights, "yolox"
            else:
                # Default to rt-detr for ultralytics models
                return ultralytics_weights, "rt-detr"
        
        # Check for Ultralytics weights in training/weights/ subdirectory (YOLOv6/YOLOX format)
        training_weights = directory / "training" / "weights" / "best.pt"
        if training_weights.exists():
            dir_name = directory.name.lower()
            if "rtdetr" in dir_name:
                return training_weights, "rt-detr"
            elif "yolov6" in dir_name:
                return training_weights, "yolov6"
            elif "yolox" in dir_name:
                return training_weights, "yolox"
            else:
                # Default to yolox for training/weights structure
                return training_weights, "yolox"
                
        # Check for direct best.pt in directory
        direct_best = directory / "best.pt"
        if direct_best.exists():
            return direct_best, "rt-detr"  # Default assumption
            
        # Check for any .pth or .pt files as fallback
        pth_files = list(directory.glob("*.pth")) + list(directory.glob("*.pt"))
        if pth_files:
            # Prefer files with "best" in name
            best_files = [f for f in pth_files if "best" in f.name.lower()]
            if best_files:
                return best_files[0], self._guess_model_type_from_path(best_files[0])
            else:
                return pth_files[0], self._guess_model_type_from_path(pth_files[0])
                
        return None, None
    
    def _guess_model_type_from_path(self, path: Path) -> str:
        """Guess model type from file path."""
        path_str = str(path).lower()
        if "rf" in path_str or "checkpoint" in path_str:
            return "rf-detr"
        elif "rtdetr" in path_str:
            return "rt-detr"
        elif "yolov6" in path_str:
            return "yolov6"
        elif "yolox" in path_str:
            return "yolox"
        else:
            return "rt-detr"  # Default
    
    def _load_model(self, weights_path: Path, model_type: str = None):
        """Load model for auto-labeling based on type."""
        try:
            import torch
            
            if model_type is None:
                model_type = self._guess_model_type_from_path(weights_path)
                
            print(f"Loading {model_type} model from {weights_path}...")
            
            # Check for GPU availability
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            print(f"Using device: {device}")
            
            if model_type == "rf-detr":
                # RF-DETR uses custom loader - try different model sizes
                from rfdetr import RFDETRBase, RFDETRMedium, RFDETRNano, RFDETRSmall
                
                # Try to determine model size from checkpoint or use nano as default
                checkpoint = torch.load(weights_path, map_location='cpu', weights_only=False)
                if 'model' in checkpoint:
                    # Training checkpoint - check the model size from the state dict
                    state_dict = checkpoint['model']
                    # Look for clues about model size in the state dict keys
                    if any('backbone.0.encoder.encoder.embeddings.position_embeddings' in key for key in state_dict.keys()):
                        # Try different model sizes to find the right one
                        models_to_try = [
                            ("nano", RFDETRNano),
                            ("small", RFDETRSmall), 
                            ("medium", RFDETRMedium),
                            ("base", RFDETRBase)
                        ]
                        
                        for size_name, model_class in models_to_try:
                            try:
                                self.model = model_class()
                                # Try loading with strict=False to handle mismatches
                                missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
                                if len(missing_keys) < len(self.model.state_dict()):  # Some keys matched
                                    print(f"✓ Loaded RF-DETR {size_name} model (with {len(missing_keys)} missing keys)")
                                    # Move to GPU if available
                                    self.model = self.model.to(device)
                                    break
                            except Exception as e:
                                continue
                        else:
                            raise Exception("Could not load checkpoint with any RF-DETR model size")
                    else:
                        # Direct weights file
                        self.model = RFDETRNano(pretrain_weights=str(weights_path))
                        self.model = self.model.to(device)
                else:
                    # Direct weights file
                    self.model = RFDETRNano(pretrain_weights=str(weights_path))
                    self.model = self.model.to(device)
            else:
                # RT-DETR, YOLOv6, YOLOX all use Ultralytics
                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))
                
                # Ultralytics models should automatically use GPU if available
                # but let's ensure they're on the right device
                if hasattr(self.model, 'to'):
                    self.model = self.model.to(device)
                    
            self.current_model_path = weights_path
            self.current_model_type = model_type
            
            # Add to available models for quick switching
            model_display = f"Custom: {weights_path.name} ({model_type.upper()})"
            existing_model = next((m for m in self.available_models if m['path'] == weights_path), None)
            if not existing_model:
                self.available_models.append({
                    "path": weights_path,
                    "type": model_type,
                    "dir": weights_path.parent.name,
                    "display": model_display
                })
            
            print("✓ Model loaded")
            return f"✓ {model_type.upper()} model loaded from {weights_path.name}"
        except Exception as e:
            error_msg = f"⚠ Could not load {model_type or 'model'}: {e}"
            print(error_msg)
            self.model = None
            self.current_model_type = None
            return error_msg
    
    def load_new_model(self, weights_path: str, model_type: str = "Auto-detect") -> str:
        """Load a new model from the GUI."""
        path = Path(weights_path)
        if not path.exists():
            return f"❌ File not found: {weights_path}"
        
        # Convert dropdown value to internal type
        if model_type == "Auto-detect":
            model_type = None
        elif model_type == "rf-detr":
            model_type = "rf-detr"
        elif model_type == "rt-detr":
            model_type = "rt-detr"
        elif model_type == "yolov6":
            model_type = "yolov6"
        elif model_type == "yolox":
            model_type = "yolox"
        
        return self._load_model(path, model_type)
    
    def load_model_from_directory(self, directory_path: str) -> str:
        """Load the best model found in a directory."""
        path = Path(directory_path)
        if not path.exists():
            return f"❌ Directory not found: {directory_path}"
        
        if not path.is_dir():
            return f"❌ Not a directory: {directory_path}"
            
        weights_path, detected_type = self.find_best_weights(path)
        if weights_path is None:
            return f"❌ No model weights found in {directory_path}"
            
        return self._load_model(weights_path, detected_type)
    
    def scan_for_models(self, return_info: bool = True) -> str:
        """Scan for available trained models in common directories."""
        runs_dir = Path("runs")
        available_models = []
        
        if runs_dir.exists():
            for subdir in runs_dir.iterdir():
                if subdir.is_dir():
                    weights_path, model_type = self.find_best_weights(subdir)
                    if weights_path:
                        available_models.append({
                            "path": weights_path,
                            "type": model_type,
                            "dir": subdir.name,
                            "display": f"{subdir.name} ({model_type.upper()})"
                        })
        
        # Store available models for quick access
        self.available_models = available_models
        
        if not return_info:
            return ""
        
        if not available_models:
            return "❌ No trained models found in 'runs/' directory"
        
        # Format as readable list
        lines = ["📂 Available Models:"]
        for i, model in enumerate(available_models, 1):
            lines.append(f"{i}. {model['dir']} -> {model['path'].name} ({model['type'].upper()})")
        
        lines.append("\n💡 Use the Model Selector dropdown above to quickly switch models")
        return "\n".join(lines)
    
    def get_available_models_list(self) -> list:
        """Get list of available models for dropdown."""
        if not self.available_models:
            self.scan_for_models(return_info=False)  # This will populate self.available_models
        
        if not self.available_models:
            return ["No models found - click '🔍 Scan for Models'"]
        
        return [model['display'] for model in self.available_models]
    
    def load_model_by_index(self, model_display: str) -> str:
        """Load a model by its display name from the available models list."""
        if not hasattr(self, 'available_models') or not self.available_models:
            return "❌ No models available. Click '🔍 Scan for Models' first."
        
        for model in self.available_models:
            if model['display'] == model_display:
                return self._load_model(model['path'], model['type'])
        
        return f"❌ Model '{model_display}' not found"
    
    def load_new_images_dir(self, images_dir: str) -> tuple[str, str, str, str]:
        """Load a new images directory from the GUI."""
        path = Path(images_dir)
        if not path.exists():
            return "<div>Directory not found</div>", "", "Image: -/-", f"❌ Directory not found: {images_dir}"
        
        if not path.is_dir():
            return "<div>Not a directory</div>", "", "Image: -/-", f"❌ Not a directory: {images_dir}"
        
        result = self._load_images(path)
        
        # Load first image
        if self.image_paths:
            img, filename = self.get_current_image()
            boxes = self.annotations.get(filename, [])
            img_html = self.generate_interactive_canvas(boxes)
            boxes_html = self._format_boxes_html(boxes)
            image_label = self._current_image_label(filename)
            status = result
            return img_html, boxes_html, image_label, status
        else:
            return "<div>No images found</div>", "", "Image: -/-", f"{result}\n⚠️ No .jpg or .png images found in directory"
    
    def get_current_model_info(self) -> str:
        """Get info about currently loaded model."""
        if self.model and self.current_model_path:
            type_info = f" ({self.current_model_type.upper()})" if self.current_model_type else ""
            return f"📦 Loaded: {self.current_model_path}{type_info}"
        elif self.model:
            return "📦 Model loaded (pretrained)"
        else:
            return "⚠️ No model loaded"
    
    def get_current_dir_info(self) -> str:
        """Get info about current images directory."""
        return f"📁 {self.images_dir} ({len(self.image_paths)} images)"

    def _current_image_label(self, filename: str) -> str:
        """Stable image index display (kept separate from status messages)."""
        if not filename or not self.image_paths:
            return "Image: -/-"
        return f"Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
    
    def get_current_image(self) -> tuple[Image.Image, str]:
        """Get current image and filename."""
        if not self.image_paths:
            return None, ""
        path = self.image_paths[self.current_idx]
        img = Image.open(path).convert("RGB")
        return img, path.name
    
    def draw_boxes_on_image(self, img: Image.Image, boxes: list[dict]) -> Image.Image:
        """Draw bounding boxes on image."""
        img_draw = img.copy()
        draw = ImageDraw.Draw(img_draw)
        
        for box in boxes:
            x1, y1, x2, y2 = box["bbox"]
            label = box.get("label", "knot")
            conf = box.get("confidence", 1.0)
            
            # Draw box
            draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
            
            # Draw corner handles for editing (small squares)
            handle_size = 6
            draw.rectangle([x1-handle_size, y1-handle_size, x1+handle_size, y1+handle_size], fill="red")
            draw.rectangle([x2-handle_size, y1-handle_size, x2+handle_size, y1+handle_size], fill="red")
            draw.rectangle([x1-handle_size, y2-handle_size, x1+handle_size, y2+handle_size], fill="red")
            draw.rectangle([x2-handle_size, y2-handle_size, x2+handle_size, y2+handle_size], fill="red")
            
            # Draw label
            text = f"{label} {conf:.2f}" if conf < 1.0 else label
            draw.text((x1, y1 - 20), text, fill="red")
        
        return img_draw
    
    def generate_interactive_canvas(self, boxes: list[dict] = None) -> str:
        """Generate HTML with interactive canvas for annotation."""
        img, filename = self.get_current_image()
        if not img:
            return "<div>No image loaded</div>"
        
        if boxes is None:
            boxes = self.annotations.get(filename, [])
        
        # Resize image for display if too large
        max_width = 1200
        max_height = 800
        if img.width > max_width or img.height > max_height:
            ratio = min(max_width / img.width, max_height / img.height)
            display_width = int(img.width * ratio)
            display_height = int(img.height * ratio)
            img = img.resize((display_width, display_height), Image.Resampling.LANCZOS)
        else:
            display_width = img.width
            display_height = img.height
        
        # Convert PIL image to base64
        import base64
        from io import BytesIO
        buffer = BytesIO()
        img.save(buffer, format="PNG", optimize=True)
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        
        # Build HTML without <script> tags (Gradio sanitizes them). The interactive
        # canvas logic runs via CANVAS_JS_ON_LOAD.
        boxes_json = json.dumps(boxes)
        boxes_escaped = html.escape(boxes_json)

        html_out = f"""
        <div style="display: inline-block; border: 1px solid #ccc; padding: 5px;">
            <div style="font-size: 12px; color: #666; margin-bottom: 4px;">Canvas Size: {display_width}x{display_height}</div>
            <textarea id="annotation-initial-boxes" style="display:none;">{boxes_escaped}</textarea>
            <div data-annotation-root="1" style="position: relative; width: {display_width}px; height: {display_height}px;">
                <img id="annotation-img" src="data:image/png;base64,{img_base64}"
                     style="position:absolute; left:0; top:0; width:{display_width}px; height:{display_height}px; z-index: 1;" />
                <canvas id="annotation-canvas" width="{display_width}" height="{display_height}"
                        style="position:absolute; left:0; top:0; width:{display_width}px; height:{display_height}px; cursor: crosshair; background: transparent; z-index: 2; pointer-events: auto;"></canvas>
            </div>
        </div>
        """
        
        return html_out
    
    def _format_boxes_html(self, boxes: list[dict]) -> str:
        """Format boxes as HTML list with delete buttons."""
        if not boxes:
            return "<div style='color: #999; font-style: italic;'>No annotations</div>"
        
        lines = ["<div style='font-family: monospace; font-size: 13px;'>"]
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box["bbox"]
            conf = box.get("confidence", 1.0)
            source = box.get("source", "manual")
            lines.append(
                f"<div style='margin: 4px 0; padding: 4px; background: #f5f5f5; border-radius: 3px; display: flex; justify-content: space-between; align-items: center;'>"
                f"<span style='flex: 1;'>{i}: [{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}] {conf:.2f} ({source})</span>"
                f"<button onclick='window.deleteBox({i})' style='background: #ff4444; color: white; border: none; padding: 2px 8px; border-radius: 3px; cursor: pointer; margin-left: 8px;'>✕</button>"
                f"</div>"
            )
        lines.append("</div>")
        return "".join(lines)

    def _parse_boxes_text(self, text: str) -> list[dict] | None:
        """Parse edited JSON from the Annotations textbox."""
        if not text:
            return None
        try:
            data = json.loads(text)
        except json.JSONDecodeError:
            return None
        if not isinstance(data, list):
            return None
        cleaned: list[dict] = []
        for item in data:
            if not isinstance(item, dict) or "bbox" not in item:
                continue
            bbox = item.get("bbox")
            if not (isinstance(bbox, list) and len(bbox) == 4):
                continue
            try:
                x1, y1, x2, y2 = [float(v) for v in bbox]
            except Exception:
                continue
            cleaned.append({
                "bbox": [x1, y1, x2, y2],
                "label": item.get("label", "knot"),
                "confidence": float(item.get("confidence", 1.0)),
                "source": item.get("source", "manual"),
            })
        return cleaned
    
    def load_image(self, direction: str = "current") -> tuple[str, str, str, str]:
        """Load image (current/next/prev)."""
        if direction == "next":
            self.current_idx = min(self.current_idx + 1, len(self.image_paths) - 1)
        elif direction == "prev":
            self.current_idx = max(self.current_idx - 1, 0)
        
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"
        
        # Load existing annotations
        boxes = self.annotations.get(filename, [])
        img_html = self.generate_interactive_canvas(boxes)
        boxes_html = self._format_boxes_html(boxes)
        image_label = self._current_image_label(filename)

        return img_html, boxes_html, image_label, ""
    
    def add_box_manual(self, x1: int, y1: int, x2: int, y2: int) -> tuple[str, str, str, str]:
        """Manually add a bounding box."""
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"
        
        # Add box
        box = {
            "bbox": [float(x1), float(y1), float(x2), float(y2)],
            "label": "knot",
            "confidence": 1.0,
            "source": "manual"
        }
        
        if filename not in self.annotations:
            self.annotations[filename] = []
        self.annotations[filename].append(box)
        self._save_annotations()
        
        # Redraw
        boxes = self.annotations[filename]
        img_html = self.generate_interactive_canvas(boxes)
        boxes_html = self._format_boxes_html(boxes)
        return img_html, boxes_html, self._current_image_label(filename), f"✓ Added box: {len(boxes)} total"
    
    def delete_last_box(self) -> tuple[str, str, str, str]:
        """Delete the last box from current image."""
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"
        
        if filename in self.annotations and self.annotations[filename]:
            self.annotations[filename].pop()
            self._save_annotations()
        
        # Redraw
        boxes = self.annotations.get(filename, [])
        img_html = self.generate_interactive_canvas(boxes)
        boxes_html = self._format_boxes_html(boxes)
        return img_html, boxes_html, self._current_image_label(filename), f"✓ Deleted last box: {len(boxes)} remaining"
    
    def delete_box_by_index(self, index: int) -> tuple[str, str, str, str]:
        """Delete a specific box by index."""
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"
        
        boxes = self.annotations.get(filename, [])
        if 0 <= index < len(boxes):
            boxes.pop(index)
            self.annotations[filename] = boxes
            self._save_annotations()
            img_html = self.generate_interactive_canvas(boxes)
            boxes_html = self._format_boxes_html(boxes)
            return img_html, boxes_html, self._current_image_label(filename), f"✓ Deleted box {index}"
        else:
            img_html = self.generate_interactive_canvas(boxes)
            boxes_html = self._format_boxes_html(boxes)
            return img_html, boxes_html, self._current_image_label(filename), "❌ Invalid box index"
    
    def clear_boxes(self) -> tuple[str, str, str, str]:
        """Clear all boxes from current image."""
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"
        
        self.annotations[filename] = []
        self._save_annotations()
        
        boxes = []
        img_html = self.generate_interactive_canvas(boxes)
        boxes_html = self._format_boxes_html(boxes)
        return img_html, boxes_html, self._current_image_label(filename), "✓ Cleared all boxes"
    
    def auto_label_current(self, threshold: float = 0.5) -> tuple[str, str, str, str]:
        """Auto-label current image using loaded model."""
        if not self.model:
            return "<div>No model loaded</div>", "", "Image: -/-", "❌ No model loaded"
        
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"
        
        try:
            # Run inference based on model type
            if self.current_model_type == "rf-detr":
                # RF-DETR custom prediction
                detections = self.model.predict(img, threshold=threshold)
                boxes = []
                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
                    boxes.append({
                        "bbox": [float(x1), float(y1), float(x2), float(y2)],
                        "label": "knot",
                        "confidence": conf,
                        "source": "auto"
                    })
            else:
                # Ultralytics models (RT-DETR, YOLOv6, YOLOX)
                results = self.model.predict(
                    source=img,
                    conf=threshold,
                    save=False,
                    verbose=False
                )
                
                boxes = []
                for result in results:
                    for box in result.boxes:
                        x1, y1, x2, y2 = box.xyxy[0].tolist()
                        conf = float(box.conf[0])
                        boxes.append({
                            "bbox": [x1, y1, x2, y2],
                            "label": "knot",
                            "confidence": conf,
                            "source": "auto"
                        })
            
            # Add to existing annotations
            if filename not in self.annotations:
                self.annotations[filename] = []
            self.annotations[filename].extend(boxes)
            self._save_annotations()
            
            # Redraw
            img_html = self.generate_interactive_canvas(self.annotations[filename])
            boxes_html = self._format_boxes_html(self.annotations[filename])
            return img_html, boxes_html, self._current_image_label(filename), f"🤖 Auto-labeled: {len(boxes)} detections"
            
        except Exception as e:
            boxes = self.annotations.get(filename, [])
            img_html = self.generate_interactive_canvas(boxes)
            return img_html, self._format_boxes_html(boxes), self._current_image_label(filename), f"❌ Auto-label failed: {e}"
    
    def save_canvas_changes(self, boxes_json: str) -> tuple[str, str, str, str]:
        """Auto-save changes made in the canvas."""
        img, filename = self.get_current_image()
        if not img:
            return "<div>No images</div>", "", "Image: -/-", "No images"

        try:
            if boxes_json:
                boxes = json.loads(boxes_json)
                self.annotations[filename] = boxes
                self._save_annotations()
                status = ""
            else:
                boxes = self.annotations.get(filename, [])
                status = ""
        except json.JSONDecodeError:
            boxes = self.annotations.get(filename, [])
            status = "❌ Invalid canvas data"

        img_html = self.generate_interactive_canvas(boxes)
        boxes_html = self._format_boxes_html(boxes)

        return img_html, boxes_html, self._current_image_label(filename), status
    
    def _save_annotations(self):
        """Save annotations to JSON file."""
        with self.ann_file.open("w") as f:
            json.dump(self.annotations, f, indent=2)
    
    def export_to_coco(self, output_path: Path):
        """Export annotations to COCO format."""
        coco_data = {
            "images": [],
            "annotations": [],
            "categories": [{"id": 0, "name": "knot", "supercategory": "defect"}]
        }
        
        ann_id = 0
        for img_id, img_path in enumerate(self.image_paths):
            filename = img_path.name
            img = Image.open(img_path)
            width, height = img.size
            
            coco_data["images"].append({
                "id": img_id,
                "file_name": filename,
                "width": width,
                "height": height
            })
            
            # Add annotations
            boxes = self.annotations.get(filename, [])
            for box in boxes:
                if isinstance(box, dict) and "bbox" in box:
                    x1, y1, x2, y2 = box["bbox"]
                    score = box.get("confidence", 1.0)
                else:
                    # Backward/experimental compatibility: [x1, y1, x2, y2]
                    x1, y1, x2, y2 = box
                    score = 1.0
                w = x2 - x1
                h = y2 - y1
                
                coco_data["annotations"].append({
                    "id": ann_id,
                    "image_id": img_id,
                    "category_id": 0,
                    "bbox": [x1, y1, w, h],
                    "area": w * h,
                    "iscrowd": 0,
                    "score": score
                })
                ann_id += 1
        
        with output_path.open("w") as f:
            json.dump(coco_data, f, indent=2)
        
        return f"✓ Exported {len(coco_data['annotations'])} annotations to {output_path}"
    
    def get_model_path_from_display(self, model_display: str) -> Path | None:
        """Get the actual model path from a display name."""
        if not hasattr(self, 'available_models') or not self.available_models:
            return None
        
        for model in self.available_models:
            if model['display'] == model_display:
                return model['path']
        
        return None


def create_ui(app: AnnotationApp) -> gr.Blocks:
    """Create Gradio UI."""
    
    with gr.Blocks(title="Knot Annotation Tool") as demo:
        gr.Markdown("""
        # Wood Knot Annotation Tool
        **Label -> Auto-Label -> Export**
        
        - Manually annotate images or use **Auto-Label** with your trained model
        - Export annotations to COCO format for training
        - Use separate training and deployment scripts for model development
        """)
        
        # Settings section at the top
        with gr.Accordion("Settings", open=False):
            with gr.Row():
                with gr.Column():
                    images_dir_input = gr.Textbox(
                        label="Images Directory",
                        value=str(app.images_dir),
                        placeholder="/path/to/images"
                    )
                    load_images_btn = gr.Button("📁 Load Images Directory")
                    dir_info = gr.Textbox(label="Current Directory", value=app.get_current_dir_info(), interactive=False)
                
                with gr.Column():
                    # Quick Model Selector
                    model_selector = gr.Dropdown(
                        choices=app.get_available_models_list(),
                        value=None,
                        label="🚀 Quick Model Switcher",
                        info="Select from available trained models (refresh with scan)",
                        allow_custom_value=True
                    )
                    quick_load_btn = gr.Button("⚡ Load Selected Model", variant="primary")
                    
                    # Manual Model Loading
                    model_weights_input = gr.Textbox(
                        label="Model Weights Path",
                        value=str(app.current_model_path) if app.current_model_path else "",
                        placeholder="runs/training/checkpoint_best_total.pth"
                    )
                    model_type_dropdown = gr.Dropdown(
                        choices=["Auto-detect", "rf-detr", "rt-detr", "yolov6", "yolox"],
                        value="Auto-detect",
                        label="Model Type",
                        info="Auto-detect will try to determine from file path"
                    )
                    with gr.Row():
                        load_model_btn = gr.Button("🤖 Load Model Weights")
                        scan_models_btn = gr.Button("🔍 Scan for Models")
                    model_info = gr.Textbox(label="Current Model", value=app.get_current_model_info(), interactive=False)
        
        with gr.Row():
            with gr.Column(scale=3):
                image_display = gr.HTML(
                    label="Current Image",
                    value="<div style='width: 800px; height: 400px; border: 1px solid #ccc; display: flex; align-items: center; justify-content: center; color: #666;'>Load images from Settings to start annotating</div>",
                    js_on_load=CANVAS_JS_ON_LOAD,
                )
                
                with gr.Row():
                    prev_btn = gr.Button("⬅️ Previous")
                    next_btn = gr.Button("Next ➡️")
                    auto_label_btn = gr.Button("🤖 Auto-Label", variant="primary")
                
                # Hidden textbox to store canvas boxes data
                canvas_boxes_data = gr.Textbox(visible=False, elem_id="canvas-boxes-data")
                
                with gr.Row():
                    threshold_slider = gr.Slider(0.1, 0.9, DEFAULT_DETECTION_THRESHOLD, label="Detection Threshold")
            
            with gr.Column(scale=1):
                image_index_text = gr.Textbox(label="Image", lines=1, interactive=False)
                info_text = gr.Textbox(label="Status", lines=2)
                boxes_html = gr.HTML(label="Annotations")
                delete_box_index = gr.Number(visible=False, value=-1)
                
                gr.Markdown("### Manual Annotation")
                with gr.Row():
                    x1_input = gr.Number(label="x1", value=100)
                    y1_input = gr.Number(label="y1", value=100)
                with gr.Row():
                    x2_input = gr.Number(label="x2", value=200)
                    y2_input = gr.Number(label="y2", value=200)
                
                add_box_btn = gr.Button("➕ Add Box")
                delete_btn = gr.Button("🗑️ Delete Last")
                clear_btn = gr.Button("❌ Clear All")
                
                gr.Markdown("### Export Annotations")
                export_path = gr.Textbox(
                    label="Export Path",
                    value="annotations_coco.json"
                )
                export_btn = gr.Button("Export COCO")
                export_result = gr.Textbox(label="Export Result", lines=1)
        
        # Event handlers
        def on_load():
            return app.load_image("current")
        
        # Settings handlers
        load_images_btn.click(
            app.load_new_images_dir,
            inputs=[images_dir_input],
            outputs=[image_display, boxes_html, image_index_text, info_text]
        ).then(
            lambda: (app.get_current_dir_info(), app.get_current_model_info()),
            outputs=[dir_info, model_info]
        )
        
        load_model_btn.click(
            app.load_new_model,
            inputs=[model_weights_input, model_type_dropdown],
            outputs=[model_info]
        ).then(
            app.get_available_models_list,
            outputs=[model_selector]
        )
        
        scan_models_btn.click(
            app.scan_for_models,
            outputs=[model_info]
        ).then(
            app.get_available_models_list,
            outputs=[model_selector]
        )
        
        quick_load_btn.click(
            app.load_model_by_index,
            inputs=[model_selector],
            outputs=[model_info]
        ).then(
            app.get_available_models_list,
            outputs=[model_selector]
        )
        
        prev_btn.click(
            lambda: app.load_image("prev"),
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        next_btn.click(
            lambda: app.load_image("next"),
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        auto_label_btn.click(
            lambda t: app.auto_label_current(t),
            inputs=[threshold_slider],
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        # Auto-save when canvas changes
        canvas_boxes_data.change(
            app.save_canvas_changes,
            inputs=[canvas_boxes_data],
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        # Delete box handler (called from HTML button clicks via JS)
        delete_box_index.change(
            lambda idx: app.delete_box_by_index(int(idx)) if idx >= 0 else (None, None, None, None),
            inputs=[delete_box_index],
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        add_box_btn.click(
            app.add_box_manual,
            inputs=[x1_input, y1_input, x2_input, y2_input],
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        delete_btn.click(
            app.delete_last_box,
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        clear_btn.click(
            app.clear_boxes,
            outputs=[image_display, boxes_html, image_index_text, info_text]
        )
        
        export_btn.click(
            lambda path: app.export_to_coco(Path(path)),
            inputs=[export_path],
            outputs=[export_result]
        )
        
        # Load first image on start
        demo.load(on_load, outputs=[image_display, boxes_html, image_index_text, info_text])
    
    return demo


def main():
    parser = argparse.ArgumentParser(description="Simple annotation GUI with auto-labeling")
    parser.add_argument(
        "--images-dir",
        type=Path,
        default=Path(DEFAULT_IMAGES_DIR) if DEFAULT_IMAGES_DIR else None,
        help="Default directory with images (can be changed in GUI)"
    )
    parser.add_argument(
        "--model-weights",
        type=Path,
        default=Path(DEFAULT_MODEL_WEIGHTS) if DEFAULT_MODEL_WEIGHTS else None,
        help="Default trained model for auto-labeling (can be changed in GUI)"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=DEFAULT_PORT,
        help="Port to run the GUI on"
    )
    args = parser.parse_args()
    
    # Validate paths if provided
    if args.images_dir and not args.images_dir.exists():
        print(f"⚠️ Warning: Images directory not found: {args.images_dir}")
        print("You can load a different directory from the GUI Settings")
        args.images_dir = None
    
    if args.model_weights and not args.model_weights.exists():
        print(f"⚠️ Warning: Model weights not found: {args.model_weights}")
        print("You can load different weights from the GUI Settings")
        args.model_weights = None
    
    # Create app
    app = AnnotationApp(args.images_dir, args.model_weights)
    
    # Scan for available models on startup
    app.scan_for_models(return_info=False)
    
    # Create and launch UI
    demo = create_ui(app)
    
    print(f"\n{'='*60}")
    print(f"🚀 Starting annotation tool...")
    if args.images_dir:
        print(f"📁 Default images: {args.images_dir} ({len(app.image_paths)} images)")
    else:
        print(f"📁 No default images - load directory from Settings")
    if app.model:
        print(f"🤖 Model: Loaded from {args.model_weights}")
    else:
        print(f"⚠️  No model loaded - load from Settings or train one")
    print(f"💡 You can change images directory and model weights from the Settings panel")
    print(f"{'='*60}\n")
    
    # Combine canvas JS with delete button handler
    combined_js = CANVAS_GLOBAL_JS + r"""
// Wire delete buttons to hidden number input
window.deleteBox = function(index) {
  const hiddenInput = document.querySelector('input[type=number][style*=display], input[type=number].\\!hidden');
  if (hiddenInput) {
    hiddenInput.value = index;
    hiddenInput.dispatchEvent(new Event('input', { bubbles: true }));
    // Reset after triggering
    setTimeout(() => { hiddenInput.value = -1; }, 100);
  }
};
"""
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=args.port,
        js=combined_js,
        share=False
    )


if __name__ == "__main__":
    main()