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 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"


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
        self.training_process = None
        self.training_thread = None
        self.training_status = "Not training"
        
        # 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[Image.Image | None, str, str]:
        """Load a new images directory from the GUI."""
        path = Path(images_dir)
        if not path.exists():
            return None, "", f"❌ Directory not found: {images_dir}"
        
        if not path.is_dir():
            return None, "", 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_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
            boxes_text = self._format_boxes_text(boxes)
            info = f"{result}\nImage 1/{len(self.image_paths)}: {filename}"
            return img_with_boxes, boxes_text, info
        else:
            return None, "", 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 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 label
            text = f"{label} {conf:.2f}" if conf < 1.0 else label
            draw.text((x1, y1 - 20), text, fill="red")
        
        return img_draw
    
    def auto_label_current(self, threshold: float = 0.5) -> tuple[Image.Image, str, str]:
        """Auto-label current image using loaded model."""
        if not self.model:
            return None, "", "❌ No model loaded"
        
        img, filename = self.get_current_image()
        if not img:
            return None, "", "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_with_boxes = self.draw_boxes_on_image(img, self.annotations[filename])
            boxes_text = self._format_boxes_text(self.annotations[filename])
            info = f"🤖 Auto-labeled: {len(boxes)} detections | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
            
            return img_with_boxes, boxes_text, info
            
        except Exception as e:
            return img, self._format_boxes_text(self.annotations.get(filename, [])), f"❌ Auto-label failed: {e}"
    
    def _format_boxes_text(self, boxes: list[dict]) -> str:
        """Format boxes for display."""
        if not boxes:
            return "No annotations"
        
        lines = []
        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"{i}: [{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}] conf={conf:.2f} ({source})")
        
        return "\n".join(lines)
    
    def load_image(self, direction: str = "current") -> tuple[Image.Image, 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 None, "", "No images"
        
        # Load existing annotations
        boxes = self.annotations.get(filename, [])
        img_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
        boxes_text = self._format_boxes_text(boxes)
        info = f"Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img_with_boxes, boxes_text, info
    
    def add_box_manual(self, x1: int, y1: int, x2: int, y2: int) -> tuple[Image.Image, str, str]:
        """Manually add a bounding box."""
        img, filename = self.get_current_image()
        if not img:
            return None, "", "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_with_boxes = self.draw_boxes_on_image(img, boxes)
        boxes_text = self._format_boxes_text(boxes)
        info = f"✓ Added box: {len(boxes)} total | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img_with_boxes, boxes_text, info
    
    def delete_last_box(self) -> tuple[Image.Image, str, str]:
        """Delete the last box from current image."""
        img, filename = self.get_current_image()
        if not img:
            return None, "", "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_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
        boxes_text = self._format_boxes_text(boxes)
        info = f"✓ Deleted last box: {len(boxes)} remaining | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img_with_boxes, boxes_text, info
    
    def clear_boxes(self) -> tuple[Image.Image, str, str]:
        """Clear all boxes from current image."""
        img, filename = self.get_current_image()
        if not img:
            return None, "", "No images"
        
        self.annotations[filename] = []
        self._save_annotations()
        
        boxes_text = "No annotations"
        info = f"✓ Cleared all boxes | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img, boxes_text, info
    
    def auto_label_current(self, threshold: float = 0.5) -> tuple[Image.Image, str, str]:
        """Auto-label current image using loaded model."""
        if not self.model:
            return None, "", "❌ No model loaded"
        
        img, filename = self.get_current_image()
        if not img:
            return None, "", "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_with_boxes = self.draw_boxes_on_image(img, self.annotations[filename])
            boxes_text = self._format_boxes_text(self.annotations[filename])
            info = f"🤖 Auto-labeled: {len(boxes)} detections | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
            
            return img_with_boxes, boxes_text, info
            
        except Exception as e:
            return img, self._format_boxes_text(self.annotations.get(filename, [])), f"❌ Auto-label failed: {e}"
    
    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:
                x1, y1, x2, y2 = box["bbox"]
                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": box.get("confidence", 1.0)
                })
                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 prepare_training_dataset(self, output_dir: Path, train_split: float = 0.8, valid_split: float = 0.1):
        """Prepare dataset in RF-DETR format (train/valid/test splits)."""
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Create splits
        import random
        annotated_images = [img for img in self.image_paths if img.name in self.annotations and self.annotations[img.name]]
        
        if len(annotated_images) < 10:
            return f"⚠️ Need at least 10 annotated images, have {len(annotated_images)}"
        
        random.shuffle(annotated_images)
        n = len(annotated_images)
        train_n = int(n * train_split)
        valid_n = int(n * valid_split)
        
        splits = {
            "train": annotated_images[:train_n],
            "valid": annotated_images[train_n:train_n + valid_n],
            "test": annotated_images[train_n + valid_n:]
        }
        
        # Create directories and copy images
        import shutil
        for split_name, split_images in splits.items():
            split_dir = output_dir / split_name
            split_dir.mkdir(exist_ok=True)
            
            # Prepare COCO JSON for this split
            coco_data = {
                "images": [],
                "annotations": [],
                "categories": [{"id": 0, "name": "knot", "supercategory": "defect"}]
            }
            
            ann_id = 0
            for img_id, img_path in enumerate(split_images):
                # Copy image
                dest = split_dir / img_path.name
                shutil.copy2(img_path, dest)
                
                # Add to COCO
                img = Image.open(img_path)
                width, height = img.size
                
                coco_data["images"].append({
                    "id": img_id,
                    "file_name": img_path.name,
                    "width": width,
                    "height": height
                })
                
                # Add annotations
                boxes = self.annotations.get(img_path.name, [])
                for box in boxes:
                    x1, y1, x2, y2 = box["bbox"]
                    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
                    })
                    ann_id += 1
            
            # Save COCO JSON
            with (split_dir / "_annotations.coco.json").open("w") as f:
                json.dump(coco_data, f, indent=2)
        
        return f"✓ Dataset prepared: {len(splits['train'])} train, {len(splits['valid'])} valid, {len(splits['test'])} test"
    
    def start_training(self, framework: str, dataset_dir: str, output_dir: str, model_size: str, 
                      epochs: int, batch_size: int, lr: float, progress=gr.Progress()):
        """Start training in background."""
        dataset_path = Path(dataset_dir)
        output_path = Path(output_dir)
        
        if not dataset_path.exists():
            return "❌ Dataset directory not found"
        
        if self.training_process and self.training_process.poll() is None:
            return "⚠️ Training already in progress"
        
        output_path.mkdir(parents=True, exist_ok=True)
        
        # Build training command based on framework
        venv_python = Path(__file__).parent / ".venv/bin/python"
        
        if framework == "RF-DETR":
            train_script = Path(__file__).parent / "train_rfdetr.py"
            # Map sizes: nano->nano, small->small, medium->medium, base->base
            size_map = {"nano": "nano", "small": "small", "medium": "medium", "base": "base"}
            model_arg = size_map.get(model_size, "medium")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--grad-accum-steps", "2",  # Default grad accum
                "--lr", str(lr)
            ]
        elif framework == "RT-DETR":
            train_script = Path(__file__).parent / "train_rtdetr.py"
            # Map sizes: nano->r18, small->r34, medium->r50, base->l
            size_map = {"nano": "rtdetr-r18", "small": "rtdetr-r34", "medium": "rtdetr-r50", "base": "rtdetr-l"}
            model_arg = size_map.get(model_size, "rtdetr-r18")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--lr", str(lr)
            ]
        elif framework == "YOLOv6":
            train_script = Path(__file__).parent / "train_yolov6.py"
            # Map sizes: nano->n, small->s, medium->m, base->l
            size_map = {"nano": "yolov6n", "small": "yolov6s", "medium": "yolov6m", "base": "yolov6l"}
            model_arg = size_map.get(model_size, "yolov6n")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--lr", str(lr)
            ]
        elif framework == "YOLOX":
            train_script = Path(__file__).parent / "train_yolox.py"
            # Map sizes: nano->nano, small->s, medium->m, base->l
            size_map = {"nano": "yolox-nano", "small": "yolox-s", "medium": "yolox-m", "base": "yolox-l"}
            model_arg = size_map.get(model_size, "yolox-nano")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--lr", str(lr)
            ]
        else:
            return f"❌ Unknown framework: {framework}"
        
        # Start training process
        log_file = output_path / "training.log"
        self.training_status = f"🚀 Starting {framework} training..."
        
        def run_training():
            try:
                with log_file.open("w") as f:
                    self.training_process = subprocess.Popen(
                        cmd,
                        stdout=f,
                        stderr=subprocess.STDOUT,
                        text=True
                    )
                    self.training_status = f"⏳ Training in progress (PID: {self.training_process.pid})"
                    self.training_process.wait()
                    
                    if self.training_process.returncode == 0:
                        self.training_status = "✅ Training completed successfully!"
                        # Reload model with new weights
                        if framework == "RF-DETR":
                            # RF-DETR uses checkpoint_best_total.pth
                            best_weights = output_path / "checkpoint_best_total.pth"
                            model_type = "rf-detr"
                        elif framework == "RT-DETR":
                            # RT-DETR uses best.pt in weights/ subdirectory (Ultralytics)
                            best_weights = output_path / "weights" / "best.pt"
                            model_type = "rt-detr"
                        elif framework == "YOLOv6":
                            best_weights = output_path / "weights" / "best.pt"
                            model_type = "yolov6"
                        elif framework == "YOLOX":
                            best_weights = output_path / "weights" / "best.pt"
                            model_type = "yolox"
                        
                        if best_weights.exists():
                            self._load_model(best_weights, model_type)
                    else:
                        self.training_status = f"❌ Training failed (exit code {self.training_process.returncode})"
            except Exception as e:
                self.training_status = f"❌ Error: {e}"
        
        self.training_thread = threading.Thread(target=run_training, daemon=True)
        self.training_thread.start()
        
        return f"✓ Training started! Check {log_file} for progress"
    
    def get_training_status(self):
        """Get current training status."""
        return self.training_status
    
    def stop_training(self):
        """Stop the training process."""
        if self.training_process and self.training_process.poll() is None:
            self.training_process.terminate()
            self.training_status = "⏹️ Training stopped by user"
            return "✓ Training process terminated"
        return "⚠️ No training in progress"
    
    def export_for_oak_d(self, model_path: str, output_dir: str = "oak_d_export", img_size: int = 640):
        """Export trained model for OAK-D camera deployment."""
        try:
            weights_path = Path(model_path)
            output_path = Path(output_dir)
            
            if not weights_path.exists():
                return "❌ Model weights not found"
            
            output_path.mkdir(parents=True, exist_ok=True)
            
            # Determine model type
            model_type = self._guess_model_type_from_path(weights_path)
            
            print(f"Exporting {model_type} model for OAK-D...")
            
            if model_type == "rf-detr":
                # RF-DETR export - use existing export_onnx.py logic
                from rfdetr import RFDETRBase
                
                model = RFDETRBase(pretrain_weights=str(weights_path))
                model.export()  # Creates output/model.onnx
                
                # Move to output directory
                onnx_source = Path("output/model.onnx")
                if onnx_source.exists():
                    onnx_dest = output_path / "rf_detr_model.onnx"
                    onnx_source.rename(onnx_dest)
                    
                    return f"✓ RF-DETR exported for OAK-D!\n📁 Output: {output_path}\n🔗 Next: Convert ONNX to blob using blobconverter.luxonis.com"
                else:
                    return "❌ ONNX export failed"
                    
            else:
                # Ultralytics models (RT-DETR, YOLOv6, YOLOX)
                if model_type == "rt-detr":
                    from ultralytics import RTDETR
                    model = RTDETR(str(weights_path))
                else:
                    from ultralytics import YOLO
                    model = YOLO(str(weights_path))
                
                # Export to ONNX
                onnx_path = model.export(
                    format="onnx",
                    imgsz=img_size,
                    simplify=True,
                    opset=11,  # OAK-compatible opset
                )
                
                # Move ONNX to output directory
                if Path(onnx_path).exists():
                    final_onnx = output_path / f"{model_type}_model.onnx"
                    Path(onnx_path).rename(final_onnx)
                    onnx_path = final_onnx
                
                # Try to export to OpenVINO if available
                try:
                    openvino_path = model.export(
                        format="openvino",
                        imgsz=img_size,
                        half=False,  # Use FP32 for better compatibility
                    )
                    
                    # Move OpenVINO files to output directory
                    if Path(openvino_path).exists():
                        import shutil
                        openvino_dir = Path(openvino_path)
                        for file in openvino_dir.glob("*"):
                            if file.is_file():
                                shutil.move(str(file), str(output_path / file.name))
                        openvino_dir.rmdir()  # Remove empty dir
                    
                    return f"✓ {model_type.upper()} exported for OAK-D!\n📁 Output: {output_path}\n🔗 Next: Convert .xml/.bin to blob using blobconverter.luxonis.com"
                    
                except Exception as e:
                    # OpenVINO not available, just return ONNX
                    return f"✓ {model_type.upper()} exported to ONNX!\n📁 Output: {output_path}\n🔗 Next: Convert ONNX to blob using blobconverter.luxonis.com\n⚠️  OpenVINO not available: {str(e)}"
                
        except Exception as e:
            return f"❌ Export failed: {str(e)}"


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 → Train → Auto-Label → Repeat**
        
        - Manually annotate images or use **Auto-Label** with your trained model
        - Export and prepare dataset for training
        - Train **RF-DETR, RT-DETR, YOLOv6, or YOLOX** (all free for commercial use!)
        - Optimized for OAK-D camera deployment
        - Use trained model to auto-label more images
        """)
        
        # 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.Image(label="Current Image", type="pil")
                
                with gr.Row():
                    prev_btn = gr.Button("⬅️ Previous")
                    next_btn = gr.Button("Next ➡️")
                    auto_label_btn = gr.Button("🤖 Auto-Label", variant="primary")
                
                with gr.Row():
                    threshold_slider = gr.Slider(0.1, 0.9, DEFAULT_DETECTION_THRESHOLD, label="Detection Threshold")
            
            with gr.Column(scale=1):
                info_text = gr.Textbox(label="Status", lines=2)
                boxes_text = gr.Textbox(label="Annotations", lines=10)
                
                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 & Training")
                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)
        
        # Training tab
        with gr.Tab("🎯 Training"):
            gr.Markdown("""
            ### Train Object Detection Model
            
            **Choose your framework:**
            - **RF-DETR** (MIT): Custom transformer, high accuracy
            - **RT-DETR** (Apache 2.0): Ultralytics transformer, great accuracy
            - **YOLOv6** (MIT): Fast, proven on OAK cameras
            - **YOLOX** (MIT): Similar to YOLOv6, slight differences
            
            **All MIT/Apache 2.0 licensed - free for commercial use!** ✅
            
            **Steps:**
            1. Annotate at least 50-100 images in the Annotation tab
            2. Click "Prepare Dataset" to create train/valid/test splits
            3. Select your framework and configure training parameters
            4. Click "Start Training" (runs in background)
            5. After training, export for OAK-D deployment
            """)
            
            with gr.Row():
                with gr.Column():
                    dataset_prep_dir = gr.Textbox(
                        label="Dataset Output Directory",
                        value="dataset_prepared"
                    )
                    train_split = gr.Slider(0.5, 0.9, 0.8, label="Train Split Ratio")
                    valid_split = gr.Slider(0.05, 0.3, 0.1, label="Valid Split Ratio")
                    prep_btn = gr.Button("📦 Prepare Dataset", variant="secondary")
                    prep_result = gr.Textbox(label="Preparation Result", lines=2)
                
                with gr.Column():
                    gr.Markdown("### Training Configuration")
                    model_framework = gr.Dropdown(
                        choices=["RF-DETR", "RT-DETR", "YOLOv6", "YOLOX"],
                        value="RT-DETR",
                        label="Model Framework",
                        info="All MIT/Apache 2.0 licensed - free for commercial use. Optimized for OAK cameras."
                    )
                    train_dataset_dir = gr.Textbox(
                        label="Dataset Directory",
                        value="dataset_prepared"
                    )
                    train_output_dir = gr.Textbox(
                        label="Output Directory",
                        value="runs/gui_training"
                    )
                    model_size = gr.Dropdown(
                        choices=["nano", "small", "medium", "base"],
                        value=DEFAULT_MODEL_SIZE,
                        label="Model Size"
                    )
                    epochs = gr.Slider(5, 100, DEFAULT_TRAIN_EPOCHS, step=5, label="Epochs")
                    batch_size = gr.Slider(1, 16, DEFAULT_BATCH_SIZE, step=1, label="Batch Size")
                    learning_rate = gr.Number(value=DEFAULT_LEARNING_RATE, label="Learning Rate")
            
            with gr.Row():
                start_train_btn = gr.Button("🚀 Start Training", variant="primary")
                stop_train_btn = gr.Button("⏹️ Stop Training", variant="stop")
                refresh_status_btn = gr.Button("🔄 Refresh Status")
            
            training_status = gr.Textbox(
                label="Training Status",
                value="Not training",
                lines=3
            )
            
            gr.Markdown("""
            **Note**: Training runs in the background. You can continue annotating while training.
            Check the training log file for detailed progress.
            """)
        
        # OAK-D Deployment tab
        with gr.Tab("🚀 OAK-D Deployment"):
            gr.Markdown("""
            ### Deploy Trained Model to OAK-D Camera
            
            Convert your trained model to work with the **OAK-D 4 Pro** camera for real-time edge inference.
            
            **Supported Models**: RF-DETR, RT-DETR, YOLOv6, YOLOX
            
            **Process**:
            1. Select a trained model from your runs/ directory
            2. Export to ONNX and OpenVINO formats
            3. Convert OpenVINO model to blob for OAK-D
            4. Deploy blob to your OAK-D camera
            """)
            
            with gr.Row():
                with gr.Column():
                    oak_model_selector = gr.Dropdown(
                        choices=app.get_available_models_list(),
                        value=None,
                        label="Select Trained Model",
                        info="Choose a model from your training runs",
                        allow_custom_value=True
                    )
                    oak_output_dir = gr.Textbox(
                        label="Output Directory",
                        value="oak_d_deployment",
                        placeholder="oak_d_deployment"
                    )
                    oak_img_size = gr.Dropdown(
                        choices=[320, 416, 512, 640, 800, 1024],
                        value=640,
                        label="Image Size",
                        info="Input size for the model (should match training)"
                    )
                    
                    with gr.Row():
                        oak_scan_btn = gr.Button("🔍 Scan for Models")
                        oak_export_btn = gr.Button("🚀 Export for OAK-D", variant="primary")
                    
                    oak_status = gr.Textbox(
                        label="Export Status",
                        value="Ready to export",
                        lines=4
                    )
                
                with gr.Column():
                    gr.Markdown("""
                    ### 📋 Deployment Instructions
                    
                    **After Export:**
                    1. **Test OpenVINO Model** (optional):
                       ```bash
                       python -c "from openvino.runtime import Core; core = Core(); model = core.read_model('model.xml'); print('✓ Model loaded')"
                       ```
                    
                    2. **Convert to Blob**:
                       - Go to: https://blobconverter.luxonis.com/
                       - Upload your `.xml` and `.bin` files
                       - Select OAK-D device
                       - Download the `.blob` file
                    
                    3. **Deploy to OAK-D**:
                       - Use DepthAI Python API
                       - Or use OAK-D examples with your blob
                    
                    ### 💡 Tips
                    - Use **FP32** for best accuracy (default)
                    - **Nano models** work best on edge devices
                    - Test inference speed vs accuracy trade-off
                    """)
        
        # 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_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_text, info_text]
        )
        
        next_btn.click(
            lambda: app.load_image("next"),
            outputs=[image_display, boxes_text, info_text]
        )
        
        auto_label_btn.click(
            lambda t: app.auto_label_current(t),
            inputs=[threshold_slider],
            outputs=[image_display, boxes_text, info_text]
        )
        
        add_box_btn.click(
            app.add_box_manual,
            inputs=[x1_input, y1_input, x2_input, y2_input],
            outputs=[image_display, boxes_text, info_text]
        )
        
        delete_btn.click(
            app.delete_last_box,
            outputs=[image_display, boxes_text, info_text]
        )
        
        clear_btn.click(
            app.clear_boxes,
            outputs=[image_display, boxes_text, info_text]
        )
        
        export_btn.click(
            lambda path: app.export_to_coco(Path(path)),
            inputs=[export_path],
            outputs=[export_result]
        )
        
        # Training handlers
        prep_btn.click(
            lambda out, train, valid: app.prepare_training_dataset(Path(out), train, valid),
            inputs=[dataset_prep_dir, train_split, valid_split],
            outputs=[prep_result]
        )
        
        start_train_btn.click(
            app.start_training,
            inputs=[model_framework, train_dataset_dir, train_output_dir, model_size, epochs, batch_size, learning_rate],
            outputs=[training_status]
        )
        
        stop_train_btn.click(
            app.stop_training,
            outputs=[training_status]
        )
        
        refresh_status_btn.click(
            app.get_training_status,
            outputs=[training_status]
        )
        
        # OAK-D Deployment handlers
        oak_scan_btn.click(
            app.scan_for_models,
            outputs=[oak_status]
        ).then(
            app.get_available_models_list,
            outputs=[oak_model_selector]
        )
        
        oak_export_btn.click(
            app.export_for_oak_d,
            inputs=[oak_model_selector, oak_output_dir, oak_img_size],
            outputs=[oak_status]
        )
        
        # Load first image on start
        demo.load(on_load, outputs=[image_display, boxes_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)"
    )
    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")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )


if __name__ == "__main__":
    main()