diff --git a/README.md b/README.md
index cbe4c95b..0025a3d1 100644
--- a/README.md
+++ b/README.md
@@ -48,6 +48,23 @@ cd saw_mill_knot_detection
python -m venv .venv
source .venv/bin/activate
+# Install dependencies
+pip install -r requirements.txt
+```
+
+### 2. Run the Annotation GUI
+
+The repository includes an automated script that handles virtual environment activation:
+
+```bash
+# Run the GUI (automatically detects and activates venv/conda environment)
+./run_gui.sh
+
+# Or run manually
+source .venv/bin/activate # or conda activate your_env
+python annotation_gui.py
+```
+
# Install dependencies
pip install -U pip
pip install ultralytics gradio rfdetr
diff --git a/annotation_gui.py b/annotation_gui.py
index d5dcd9b4..f03768f6 100644
--- a/annotation_gui.py
+++ b/annotation_gui.py
@@ -328,14 +328,14 @@ class AnnotationApp:
return f"❌ Model '{model_display}' not found"
- def load_new_images_dir(self, images_dir: str) -> tuple[Image.Image | None, str, str]:
+ def load_new_images_dir(self, images_dir: str) -> tuple[str, 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}"
+ return "
Directory not found
", "", f"❌ Directory not found: {images_dir}"
if not path.is_dir():
- return None, "", f"❌ Not a directory: {images_dir}"
+ return "Not a directory
", "", f"❌ Not a directory: {images_dir}"
result = self._load_images(path)
@@ -343,12 +343,12 @@ class AnnotationApp:
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
+ img_html = self.generate_interactive_canvas(boxes)
boxes_text = self._format_boxes_text(boxes)
info = f"{result}\nImage 1/{len(self.image_paths)}: {filename}"
- return img_with_boxes, boxes_text, info
+ return img_html, boxes_text, info
else:
- return None, "", f"{result}\n⚠️ No .jpg or .png images found in directory"
+ return "No images found
", "", f"{result}\n⚠️ No .jpg or .png images found in directory"
def get_current_model_info(self) -> str:
"""Get info about currently loaded model."""
@@ -385,73 +385,285 @@ class AnnotationApp:
# Draw box
draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
+ # Draw corner handles for editing (small squares)
+ handle_size = 6
+ draw.rectangle([x1-handle_size, y1-handle_size, x1+handle_size, y1+handle_size], fill="red")
+ draw.rectangle([x2-handle_size, y1-handle_size, x2+handle_size, y1+handle_size], fill="red")
+ draw.rectangle([x1-handle_size, y2-handle_size, x1+handle_size, y2+handle_size], fill="red")
+ draw.rectangle([x2-handle_size, y2-handle_size, x2+handle_size, y2+handle_size], fill="red")
+
# 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"
-
+ def generate_interactive_canvas(self, boxes: list[dict] = None) -> str:
+ """Generate HTML with interactive canvas for annotation."""
img, filename = self.get_current_image()
if not img:
- return None, "", "No images"
+ return "No image loaded
"
- 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
- )
+ if boxes is None:
+ boxes = self.annotations.get(filename, [])
+
+ # Resize image for display if too large
+ max_width = 1200
+ max_height = 800
+ if img.width > max_width or img.height > max_height:
+ ratio = min(max_width / img.width, max_height / img.height)
+ display_width = int(img.width * ratio)
+ display_height = int(img.height * ratio)
+ img = img.resize((display_width, display_height), Image.Resampling.LANCZOS)
+ else:
+ display_width = img.width
+ display_height = img.height
+
+ # Convert PIL image to base64
+ import base64
+ from io import BytesIO
+ buffer = BytesIO()
+ img.save(buffer, format="PNG", optimize=True)
+ img_base64 = base64.b64encode(buffer.getvalue()).decode()
+
+ # Generate HTML with canvas and JavaScript
+ boxes_json = json.dumps(boxes)
+
+ html = f"""
+
+
Canvas Size: {display_width}x{display_height}
+
+
+ Canvas ready - testing image loading
+
+
+
+ """
+
+ return html
def _format_boxes_text(self, boxes: list[dict]) -> str:
"""Format boxes for display."""
@@ -467,7 +679,7 @@ class AnnotationApp:
return "\n".join(lines)
- def load_image(self, direction: str = "current") -> tuple[Image.Image, str, str]:
+ def load_image(self, direction: str = "current") -> tuple[str, str, str]:
"""Load image (current/next/prev)."""
if direction == "next":
self.current_idx = min(self.current_idx + 1, len(self.image_paths) - 1)
@@ -476,21 +688,21 @@ class AnnotationApp:
img, filename = self.get_current_image()
if not img:
- return None, "", "No images"
+ return "No images
", "", "No images"
# Load existing annotations
boxes = self.annotations.get(filename, [])
- img_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
+ img_html = self.generate_interactive_canvas(boxes)
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
+ return img_html, boxes_text, info
- def add_box_manual(self, x1: int, y1: int, x2: int, y2: int) -> tuple[Image.Image, str, str]:
+ def add_box_manual(self, x1: int, y1: int, x2: int, y2: int) -> tuple[str, str, str]:
"""Manually add a bounding box."""
img, filename = self.get_current_image()
if not img:
- return None, "", "No images"
+ return "No images
", "", "No images"
# Add box
box = {
@@ -507,17 +719,17 @@ class AnnotationApp:
# Redraw
boxes = self.annotations[filename]
- img_with_boxes = self.draw_boxes_on_image(img, boxes)
+ img_html = self.generate_interactive_canvas(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
+ return img_html, boxes_text, info
- def delete_last_box(self) -> tuple[Image.Image, str, str]:
+ def delete_last_box(self) -> tuple[str, str, str]:
"""Delete the last box from current image."""
img, filename = self.get_current_image()
if not img:
- return None, "", "No images"
+ return "No images
", "", "No images"
if filename in self.annotations and self.annotations[filename]:
self.annotations[filename].pop()
@@ -525,34 +737,36 @@ class AnnotationApp:
# Redraw
boxes = self.annotations.get(filename, [])
- img_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
+ img_html = self.generate_interactive_canvas(boxes)
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
+ return img_html, boxes_text, info
- def clear_boxes(self) -> tuple[Image.Image, str, str]:
+ def clear_boxes(self) -> tuple[str, str, str]:
"""Clear all boxes from current image."""
img, filename = self.get_current_image()
if not img:
- return None, "", "No images"
+ return "No images
", "", "No images"
self.annotations[filename] = []
self._save_annotations()
+ boxes = []
+ img_html = self.generate_interactive_canvas(boxes)
boxes_text = "No annotations"
info = f"✓ Cleared all boxes | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
- return img, boxes_text, info
+ return img_html, boxes_text, info
- def auto_label_current(self, threshold: float = 0.5) -> tuple[Image.Image, str, str]:
+ def auto_label_current(self, threshold: float = 0.5) -> tuple[str, str, str]:
"""Auto-label current image using loaded model."""
if not self.model:
- return None, "", "❌ No model loaded"
+ return "No model loaded
", "", "❌ No model loaded"
img, filename = self.get_current_image()
if not img:
- return None, "", "No images"
+ return "No images
", "", "No images"
try:
# Run inference based on model type
@@ -598,14 +812,41 @@ class AnnotationApp:
self._save_annotations()
# Redraw
- img_with_boxes = self.draw_boxes_on_image(img, self.annotations[filename])
+ img_html = self.generate_interactive_canvas(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
+ return img_html, boxes_text, info
except Exception as e:
- return img, self._format_boxes_text(self.annotations.get(filename, [])), f"❌ Auto-label failed: {e}"
+ boxes = self.annotations.get(filename, [])
+ img_html = self.generate_interactive_canvas(boxes)
+ return img_html, self._format_boxes_text(boxes), f"❌ Auto-label failed: {e}"
+
+ def save_canvas_changes(self, boxes_json: str) -> tuple[str, str, str]:
+ """Save changes made in the interactive canvas."""
+ img, filename = self.get_current_image()
+ if not img:
+ return "No images
", "", "No images"
+
+ try:
+ # Parse the boxes from JSON
+ if boxes_json:
+ boxes = json.loads(boxes_json)
+ self.annotations[filename] = boxes
+ self._save_annotations()
+ info = f"✓ Saved {len(boxes)} boxes | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
+ else:
+ boxes = self.annotations.get(filename, [])
+ info = f"✓ No changes to save | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
+ except json.JSONDecodeError:
+ boxes = self.annotations.get(filename, [])
+ info = f"❌ Invalid boxes data | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
+
+ img_html = self.generate_interactive_canvas(boxes)
+ boxes_text = self._format_boxes_text(boxes)
+
+ return img_html, boxes_text, info
def _save_annotations(self):
"""Save annotations to JSON file."""
@@ -1041,12 +1282,16 @@ def create_ui(app: AnnotationApp) -> gr.Blocks:
with gr.Row():
with gr.Column(scale=3):
- image_display = gr.Image(label="Current Image", type="pil")
+ image_display = gr.HTML(label="Current Image", value="Load images from Settings to start annotating
")
with gr.Row():
prev_btn = gr.Button("⬅️ Previous")
next_btn = gr.Button("Next ➡️")
auto_label_btn = gr.Button("🤖 Auto-Label", variant="primary")
+ save_canvas_btn = gr.Button("💾 Save Canvas Changes")
+
+ # Hidden textbox to store canvas boxes data
+ canvas_boxes_data = gr.Textbox(visible=False)
with gr.Row():
threshold_slider = gr.Slider(0.1, 0.9, DEFAULT_DETECTION_THRESHOLD, label="Detection Threshold")
@@ -1276,6 +1521,19 @@ def create_ui(app: AnnotationApp) -> gr.Blocks:
outputs=[image_display, boxes_text, info_text]
)
+ save_canvas_btn.click(
+ app.save_canvas_changes,
+ inputs=[canvas_boxes_data],
+ outputs=[image_display, boxes_text, info_text],
+ js="""() => {
+ const hiddenInput = document.getElementById('canvas-boxes-data');
+ if (hiddenInput) {
+ return hiddenInput.value;
+ }
+ return '';
+ }"""
+ )
+
add_box_btn.click(
app.add_box_manual,
inputs=[x1_input, y1_input, x2_input, y2_input],