From e61a4a5f81e644ee53081c7f40f37b24dbec143e Mon Sep 17 00:00:00 2001 From: dillonj Date: Tue, 23 Dec 2025 16:17:19 -0700 Subject: [PATCH] run gui script --- README.md | 17 ++ annotation_gui.py | 420 +++++++++++++++++++++++++++++++++++++--------- 2 files changed, 356 insertions(+), 81 deletions(-) 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],