From d3664693a8267245660e896cc5b14b29b44d4002 Mon Sep 17 00:00:00 2001 From: dillonj Date: Mon, 22 Dec 2025 17:56:59 -0700 Subject: [PATCH] step working on ease of deployment for oak d --- annotation_gui.py | 662 ++++++++++++++++++++++++++++++++++++---- config.py | 4 +- fix_class_ids.py | 51 ++++ fix_coco_annotations.py | 50 +++ 4 files changed, 710 insertions(+), 57 deletions(-) create mode 100644 fix_class_ids.py create mode 100644 fix_coco_annotations.py diff --git a/annotation_gui.py b/annotation_gui.py index 33fedf2b..6f7cf189 100644 --- a/annotation_gui.py +++ b/annotation_gui.py @@ -41,6 +41,8 @@ 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 @@ -74,28 +76,257 @@ class AnnotationApp: 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).""" + 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: - from ultralytics import YOLO - print(f"Loading model from {weights_path}...") - self.model = YOLO(str(weights_path)) + 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 loaded from {weights_path.name}" + return f"✓ {model_type.upper()} model loaded from {weights_path.name}" except Exception as e: - error_msg = f"⚠ Could not load model: {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) -> str: + 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}" - return self._load_model(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.""" @@ -122,7 +353,8 @@ class AnnotationApp: 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}" + 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: @@ -160,51 +392,66 @@ class AnnotationApp: return img_draw def auto_label_current(self, threshold: float = 0.5) -> tuple[Image.Image, str, str]: - """Auto-label current image with model.""" + """Auto-label current image using loaded 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 + return None, "", "❌ No model loaded" 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}" - + 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": xyxy, - "label": label, + "bbox": [float(x1), float(y1), float(x2), float(y2)], + "label": "knot", "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 + 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.""" @@ -298,6 +545,68 @@ class AnnotationApp: 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: @@ -439,7 +748,24 @@ class AnnotationApp: # Build training command based on framework venv_python = Path(__file__).parent / ".venv/bin/python" - if framework == "RT-DETR": + 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"} @@ -509,9 +835,23 @@ class AnnotationApp: 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 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) + self._load_model(best_weights, model_type) else: self.training_status = f"❌ Training failed (exit code {self.training_process.returncode})" except Exception as e: @@ -533,6 +873,88 @@ class AnnotationApp: 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: @@ -545,7 +967,7 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: - 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!) + - 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 """) @@ -563,12 +985,31 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: 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" ) - load_model_btn = gr.Button("🤖 Load Model Weights") + 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(): @@ -613,7 +1054,8 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: ### Train Object Detection Model **Choose your framework:** - - **RT-DETR** (Apache 2.0): Modern transformer, great accuracy + - **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 @@ -641,7 +1083,7 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: with gr.Column(): gr.Markdown("### Training Configuration") model_framework = gr.Dropdown( - choices=["RT-DETR", "YOLOv6", "YOLOX"], + 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." @@ -679,6 +1121,79 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: 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") @@ -695,8 +1210,28 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: load_model_btn.click( app.load_new_model, - inputs=[model_weights_input], + 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( @@ -760,6 +1295,21 @@ def create_ui(app: AnnotationApp) -> gr.Blocks: 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]) @@ -780,7 +1330,6 @@ def main(): 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 @@ -797,6 +1346,9 @@ def main(): # 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) @@ -815,7 +1367,7 @@ def main(): demo.launch( server_name="0.0.0.0", - server_port=args.port, + server_port=7860, share=False ) diff --git a/config.py b/config.py index fbba5d68..494b4a58 100644 --- a/config.py +++ b/config.py @@ -5,10 +5,10 @@ DEFAULT_IMAGES_DIR = "IMAGE/" # Directory containing wood defect images DEFAULT_MODEL_WEIGHTS = "runs/yolox_training/training/weights/best.pt" # Trained YOLOX model # Training defaults -DEFAULT_TRAIN_EPOCHS = 20 +DEFAULT_TRAIN_EPOCHS = 1 DEFAULT_BATCH_SIZE = 4 DEFAULT_LEARNING_RATE = 1e-4 -DEFAULT_MODEL_SIZE = "small" # nano, small, medium, base +DEFAULT_MODEL_SIZE = "nano" # nano, small, medium, base # Dataset split ratios DEFAULT_TRAIN_SPLIT = 0.8 diff --git a/fix_class_ids.py b/fix_class_ids.py new file mode 100644 index 00000000..083f358f --- /dev/null +++ b/fix_class_ids.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +""" +Fix COCO annotation class IDs to be 0-based instead of 1-based. +RF-DETR expects classes to start from 0 (background), but our annotations use 1-10. +""" + +import json +import sys +from pathlib import Path + +def fix_class_ids(coco_data): + """Convert category IDs and annotation category_ids from 1-based to 0-based.""" + # Fix categories: change id from 1-10 to 0-9 + for cat in coco_data['categories']: + cat['id'] -= 1 + + # Fix annotations: change category_id from 1-10 to 0-9 + for ann in coco_data['annotations']: + ann['category_id'] -= 1 + + return coco_data + +def main(): + dataset_dir = Path('dataset_coco') + + for split in ['train', 'valid', 'test']: + ann_file = dataset_dir / split / '_annotations.coco.json' + + if not ann_file.exists(): + print(f"Warning: {ann_file} not found, skipping") + continue + + print(f"Processing {ann_file}...") + + # Load data + with open(ann_file, 'r') as f: + data = json.load(f) + + # Fix class IDs + fixed_data = fix_class_ids(data) + + # Save back + with open(ann_file, 'w') as f: + json.dump(fixed_data, f, indent=2) + + print(f"✅ Fixed {len(data['annotations'])} annotations and {len(data['categories'])} categories") + + print("🎉 All COCO annotations fixed for 0-based class indexing!") + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/fix_coco_annotations.py b/fix_coco_annotations.py new file mode 100644 index 00000000..378c8e29 --- /dev/null +++ b/fix_coco_annotations.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +""" +Fix COCO annotations for RF-DETR compatibility by adding missing 'supercategory' field. +""" + +import json +from pathlib import Path + + +def fix_coco_annotations(coco_file: Path): + """Add supercategory field to categories in COCO annotation file.""" + print(f"Fixing {coco_file}...") + + # Load the COCO data + with coco_file.open('r') as f: + data = json.load(f) + + # Add supercategory field to all categories + for category in data['categories']: + if 'supercategory' not in category: + category['supercategory'] = 'wood_defect' # Default supercategory + + # Save the fixed data + with coco_file.open('w') as f: + json.dump(data, f, indent=2) + + print(f"✅ Fixed {len(data['categories'])} categories in {coco_file}") + + +def main(): + """Fix all COCO annotation files.""" + print("Fixing COCO annotations for RF-DETR compatibility...") + + coco_files = [ + Path("dataset_coco/train/_annotations.coco.json"), + Path("dataset_coco/valid/_annotations.coco.json"), + Path("dataset_coco/test/_annotations.coco.json") + ] + + for coco_file in coco_files: + if coco_file.exists(): + fix_coco_annotations(coco_file) + else: + print(f"⚠️ {coco_file} not found") + + print("\n🎉 All COCO annotations fixed for RF-DETR compatibility!") + + +if __name__ == "__main__": + main() \ No newline at end of file