""" 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 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 export_for_oak_d(self, model_display: str, output_dir: str = "oak_d_export", img_size: int = 640): """Export trained model for OAK-D camera deployment.""" try: # Convert display name to actual path weights_path = self.get_model_path_from_display(model_display) if not weights_path: return f"❌ Model '{model_display}' not found. Try clicking '🔍 Scan for Models' first." output_path = Path(output_dir) if not weights_path.exists(): return f"❌ Model weights not found at: {weights_path}" 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 import shutil docker_hint = "" if shutil.which("docker") is None: docker_hint = "\n⚠️ Docker not found (needed for offline conversion via ModelConverter)." return ( f"✓ {model_type.upper()} exported to ONNX!\n" f"📁 Output: {output_path}\n" f"🔗 Next: Convert ONNX → RVC using HubAI (online) or ModelConverter (offline).\n" f"Docs: https://docs.luxonis.com/software-v3/ai-inference/conversion/\n" f"💡 Offline conversion: Use Luxonis ModelConverter with Docker\n" f"⚠️ OpenVINO export not available: {str(e)}" f"{docker_hint}" ) 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 RVC compiled format** (recommended by Luxonis): - Online: HubAI conversion (fastest setup) - Offline: ModelConverter (requires Docker) - Docs: https://docs.luxonis.com/software-v3/ai-inference/conversion/ 3. **Deploy to OAK-D**: - Use DepthAI Python API - Or use OAK-D examples with your blob ### 💡 Tips - **Nano models** work best on edge devices - If you quantize, use real calibration images for best accuracy - 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()