Initial commit: Wood knot detection model and GUI

annotation_gui.py

@@ -0,0 +1,824 @@
+"""
+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.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 _load_model(self, weights_path: Path):
+        """Load YOLO/YOLOX model for auto-labeling (Ultralytics format)."""
+        try:
+            from ultralytics import YOLO
+            print(f"Loading model from {weights_path}...")
+            self.model = YOLO(str(weights_path))
+            self.current_model_path = weights_path
+            print("✓ Model loaded")
+            return f"✓ Model loaded from {weights_path.name}"
+        except Exception as e:
+            error_msg = f"⚠ Could not load model: {e}"
+            print(error_msg)
+            self.model = None
+            return error_msg
+    
+    def load_new_model(self, weights_path: str) -> str:
+        """Load a new model from the GUI."""
+        path = Path(weights_path)
+        if not path.exists():
+            return f"❌ File not found: {weights_path}"
+        
+        return self._load_model(path)
+    
+    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:
+            return f"📦 Loaded: {self.current_model_path}"
+        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 with model."""
+        if not self.model:
+            img, filename = self.get_current_image()
+            info = f"⚠ No model loaded | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
+            return img, "", info
+        
+        img, filename = self.get_current_image()
+        if not img:
+            return None, "", "No images"
+        
+        # Run inference with Ultralytics YOLO
+        results = self.model.predict(img, conf=threshold, verbose=False)
+        
+        # Convert to our format
+        boxes = []
+        if len(results) > 0:
+            result = results[0]  # First image result
+            if result.boxes is not None and len(result.boxes) > 0:
+                for box in result.boxes:
+                    xyxy = box.xyxy[0].cpu().numpy().tolist()  # [x1, y1, x2, y2]
+                    conf = float(box.conf[0].cpu().numpy())
+                    cls = int(box.cls[0].cpu().numpy())
+                    
+                    # Get class name if available
+                    label = result.names.get(cls, f"class_{cls}") if hasattr(result, 'names') else f"class_{cls}"
+                    
+                    boxes.append({
+                        "bbox": xyxy,
+                        "label": label,
+                        "confidence": conf,
+                        "source": "auto"
+                    })
+        
+        # Save
+        self.annotations[filename] = boxes
+        self._save_annotations()
+        
+        # Draw boxes on image
+        img_with_boxes = self.draw_boxes_on_image(img, boxes)
+        
+        # Info with image index
+        info = f"✓ Auto-labeled: {len(boxes)} boxes detected | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
+        boxes_text = self._format_boxes_text(boxes)
+        
+        return img_with_boxes, boxes_text, info
+    
+    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 _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 == "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
+                        best_weights = output_path / "checkpoint_best_total.pth"
+                        if best_weights.exists():
+                            self._load_model(best_weights)
+                    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 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 **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():
+                    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"
+                    )
+                    load_model_btn = gr.Button("🤖 Load Model Weights")
+                    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:**
+            - **RT-DETR** (Apache 2.0): Modern 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=["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.
+            """)
+        
+        # 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],
+            outputs=[model_info]
+        )
+        
+        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]
+        )
+        
+        # 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)"
+    )
+    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Port for web interface")
+    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)
+    
+    # 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=args.port,
+        share=False
+    )
+
+
+if __name__ == "__main__":
+    main()