709 lines
26 KiB
Plaintext
709 lines
26 KiB
Plaintext
|
"""
|
|||
|
|
Simple customizable annotation GUI with auto-labeling support.
|
|||
|
|
|
|||
|
|
Built with Gradio - easy to modify and extend.
|
|||
|
|
Run: python annotation_gui.py --images-dir /path/to/images
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
import numpy as np
|
|||
|
|
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 RF-DETR model for auto-labeling."""
|
|||
|
|
try:
|
|||
|
|
from rfdetr import RFDETRBase
|
|||
|
|
print(f"Loading model from {weights_path}...")
|
|||
|
|
self.model = RFDETRBase(pretrain_weights=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:
|
|||
|
|
return self.get_current_image()[0], "", "⚠ No model loaded"
|
|||
|
|
|
|||
|
|
img, filename = self.get_current_image()
|
|||
|
|
if not img:
|
|||
|
|
return None, "", "No images"
|
|||
|
|
|
|||
|
|
# Run inference
|
|||
|
|
detections = self.model.predict(img, threshold=threshold)
|
|||
|
|
|
|||
|
|
# Convert to our format
|
|||
|
|
boxes = []
|
|||
|
|
for i in range(len(detections)):
|
|||
|
|
xyxy = detections.xyxy[i].tolist()
|
|||
|
|
conf = float(detections.confidence[i]) if detections.confidence is not None else 1.0
|
|||
|
|
boxes.append({
|
|||
|
|
"bbox": xyxy,
|
|||
|
|
"label": "knot",
|
|||
|
|
"confidence": conf,
|
|||
|
|
"source": "auto"
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# Save
|
|||
|
|
self.annotations[filename] = boxes
|
|||
|
|
self._save_annotations()
|
|||
|
|
|
|||
|
|
# Draw
|
|||
|
|
img_with_boxes = self.draw_boxes_on_image(img, boxes)
|
|||
|
|
info = f"✓ Auto-labeled: {len(boxes)} boxes detected"
|
|||
|
|
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"
|
|||
|
|
|
|||
|
|
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"
|
|||
|
|
|
|||
|
|
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 = "✓ Cleared all boxes"
|
|||
|
|
|
|||
|
|
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:
|
|||
|
|
|
|||
|
|
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"
|
|||
|
|
|
|||
|
|
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.
|
|||
|
|
""
|
|||
|
|
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, 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
|
|||
|
|
venv_python = Path(__file__).parent / ".venv/bin/python"
|
|||
|
|
train_script = Path(__file__).parent / "train_rfdetr.py"
|
|||
|
|
|
|||
|
|
cmd = [
|
|||
|
|
str(venv_python),
|
|||
|
|
str(train_script),
|
|||
|
|
"--dataset-dir", str(dataset_path),
|
|||
|
|
"--output-dir", str(output_path),
|
|||
|
|
"--model", model_size,
|
|||
|
|
"--epochs", str(epochs),
|
|||
|
|
"--batch-size", str(batch_size),
|
|||
|
|
"--grad-accum-steps", str(max(1, 16 // batch_size)),
|
|||
|
|
"--lr", str(lr)
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# Start training process
|
|||
|
|
log_file = output_path / "training.log"
|
|||
|
|
self.training_status = f"🚀 Starting 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"
|
|||
|
|
json.dump(coco_data, f, indent=2)
|
|||
|
|
|
|||
|
|
return f"✓ Exported {len(coco_data['annotations'])} annotations to {output_path}"
|
|||
|
|
|
|||
|
|
|
|||
|
|
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 directly from this GUI
|
|||
|
|
- Use trained model to auto-label more images
|
|||
|
|
""")
|
|||
|
|
|
|||
|
|
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():
|
|||
|
|
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=[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]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 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")
|
|||
|
|
export_path = gr.Textbox(
|
|||
|
|
label="Output Path",
|
|||
|
|
value="annotations_coco.json"
|
|||
|
|
)
|
|||
|
|
export_btn = gr.Button("💾 Export COCO")
|
|||
|
|
export_result = gr.Textbox(label="Export Result")
|
|||
|
|
|
|||
|
|
# 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]
|
|||
|
|
)help="Default directory with images (can be changed in GUI)")
|
|||
|
|
parser.add_argument("--model-weights", type=Path, help="Default trained model for auto-labeling (can be changed in GUI)")
|
|||
|
|
parser.add_argument("--port", type=int, default=7860, 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
|
|||
|
|
|
|||
|
|
# 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, required=True, help="Directory with images")
|
|||
|
|
parser.add_argument("--model-weights", type=Path, help="Optional: trained model for auto-labeling")
|
|||
|
|
parser.add_argument("--port", type=int, default=7860, help="Port for web interface")
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
if not args.images_dir.exists():
|
|||
|
|
raise SystemExit(f"Images directory not found: {args.images_dir}")
|
|||
|
|
|
|||
|
|
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"\n{'='*60}")
|
|||
|
|
print(f"🚀 Starting annotation tool...")
|
|||
|
|
print(f"📁 Images: {args.images_dir} ({len(app.image_paths)} images)")
|
|||
|
|
if app.model:
|
|||
|
|
print(f"🤖 Model: Loaded from {args.model_weights}")
|
|||
|
|
else:
|
|||
|
|
print(f"⚠️ No model loaded (manual annotation only)")
|
|||
|
|
print(f"{'='*60}\n")
|
|||
|
|
|
|||
|
|
demo.launch(
|
|||
|
|
server_name="0.0.0.0",
|
|||
|
|
server_port=args.port,
|
|||
|
|
share=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|