dillon_stuff/saw_mill_knot_detection
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

run gui

Browse Source
This commit is contained in:
dillonj
2025-12-23 16:55:59 -07:00
parent e61a4a5f81
commit 0072a86591
2 changed files with 330 additions and 236 deletions
Download Patch File Download Diff File Expand all files Collapse all files

470
annotation_gui.py
View File

@ -10,6 +10,7 @@ To set default paths, edit config.py
from __future__ import annotations
import argparse
import html
import json
import subprocess
import threading
@ -37,6 +38,212 @@ except ImportError:
DEFAULT_MODEL_SIZE = "small"
# Gradio 6 sanitizes <script> tags inside gr.HTML content, so any canvas drawing code
# must live in the component's supported js_on_load hook.
CANVAS_JS_ON_LOAD = r"""
(() => {
const root = element;
if (!root) return;
const canvas = root.querySelector('#annotation-canvas');
const imgEl = root.querySelector('#annotation-img');
const initialBoxesEl = root.querySelector('#annotation-initial-boxes');
if (!canvas || !imgEl || !initialBoxesEl) return;
const ctx = canvas.getContext('2d');
const displayWidth = canvas.width;
const displayHeight = canvas.height;
let boxes = [];
try {
const raw = initialBoxesEl.value || initialBoxesEl.textContent || '[]';
boxes = JSON.parse(raw);
if (!Array.isArray(boxes)) boxes = [];
} catch (_) {
boxes = [];
}
const hiddenInput = document.getElementById('canvas-boxes-data');
const syncHidden = () => {
if (!hiddenInput) return;
hiddenInput.value = JSON.stringify(boxes);
hiddenInput.dispatchEvent(new Event('input', { bubbles: true }));
};
syncHidden();
let isDragging = false;
let dragStart = null;
let selectedCorner = null;
let selectedBoxIndex = -1;
let creatingBox = false;
let createStart = null;
function redraw() {
ctx.clearRect(0, 0, displayWidth, displayHeight);
// Base image is rendered via <img> below the canvas.
boxes.forEach((box) => {
const [x1, y1, x2, y2] = box.bbox;
const label = box.label || 'knot';
const conf = box.confidence || 1.0;
ctx.strokeStyle = 'red';
ctx.lineWidth = 3;
ctx.strokeRect(x1, y1, x2 - x1, y2 - y1);
const handleSize = 6;
ctx.fillStyle = 'red';
ctx.fillRect(x1 - handleSize, y1 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillRect(x2 - handleSize, y1 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillRect(x1 - handleSize, y2 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillRect(x2 - handleSize, y2 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillStyle = 'red';
ctx.font = '16px Arial';
const text = conf < 1.0 ? `${label} ${conf.toFixed(2)}` : label;
ctx.fillText(text, x1, y1 - 5);
});
if (creatingBox && createStart && dragStart) {
ctx.strokeStyle = 'blue';
ctx.lineWidth = 2;
ctx.setLineDash([5, 5]);
const x = Math.min(createStart.x, dragStart.x);
const y = Math.min(createStart.y, dragStart.y);
const w = Math.abs(createStart.x - dragStart.x);
const h = Math.abs(createStart.y - dragStart.y);
ctx.strokeRect(x, y, w, h);
ctx.setLineDash([]);
}
}
function getCornerAt(x, y) {
const handleSize = 6;
for (let i = 0; i < boxes.length; i++) {
const [x1, y1, x2, y2] = boxes[i].bbox;
const corners = [
{ x: x1, y: y1, type: 'top-left' },
{ x: x2, y: y1, type: 'top-right' },
{ x: x1, y: y2, type: 'bottom-left' },
{ x: x2, y: y2, type: 'bottom-right' },
];
for (const corner of corners) {
if (
x >= corner.x - handleSize &&
x <= corner.x + handleSize &&
y >= corner.y - handleSize &&
y <= corner.y + handleSize
) {
return { boxIndex: i, corner: corner.type, pos: corner };
}
}
}
return null;
}
// Ensure we don't double-bind if Gradio reuses the DOM node.
if (canvas.dataset.bound === '1') {
redraw();
return;
}
canvas.dataset.bound = '1';
// If the <img> fails to load, draw a message on the canvas.
imgEl.addEventListener('error', () => {
ctx.clearRect(0, 0, displayWidth, displayHeight);
ctx.fillStyle = '#ffcccc';
ctx.fillRect(0, 0, displayWidth, displayHeight);
ctx.fillStyle = 'black';
ctx.font = '16px Arial';
ctx.fillText('Image failed to load', 10, 30);
}, { once: true });
// Initial draw
redraw();
canvas.addEventListener('mousedown', (e) => {
const rect = canvas.getBoundingClientRect();
const x = (e.clientX - rect.left) * (displayWidth / rect.width);
const y = (e.clientY - rect.top) * (displayHeight / rect.height);
selectedCorner = getCornerAt(x, y);
if (selectedCorner) {
isDragging = true;
selectedBoxIndex = selectedCorner.boxIndex;
canvas.style.cursor = 'move';
} else {
creatingBox = true;
createStart = { x, y };
dragStart = { x, y };
canvas.style.cursor = 'crosshair';
}
});
canvas.addEventListener('mousemove', (e) => {
const rect = canvas.getBoundingClientRect();
const x = (e.clientX - rect.left) * (displayWidth / rect.width);
const y = (e.clientY - rect.top) * (displayHeight / rect.height);
if (isDragging && selectedCorner) {
const box = boxes[selectedBoxIndex];
if (selectedCorner.corner === 'top-left') {
box.bbox[0] = Math.min(x, box.bbox[2] - 10);
box.bbox[1] = Math.min(y, box.bbox[3] - 10);
} else if (selectedCorner.corner === 'top-right') {
box.bbox[2] = Math.max(x, box.bbox[0] + 10);
box.bbox[1] = Math.min(y, box.bbox[3] - 10);
} else if (selectedCorner.corner === 'bottom-left') {
box.bbox[0] = Math.min(x, box.bbox[2] - 10);
box.bbox[3] = Math.max(y, box.bbox[1] + 10);
} else if (selectedCorner.corner === 'bottom-right') {
box.bbox[2] = Math.max(x, box.bbox[0] + 10);
box.bbox[3] = Math.max(y, box.bbox[1] + 10);
}
syncHidden();
redraw();
return;
}
if (creatingBox && createStart) {
dragStart = { x, y };
redraw();
return;
}
const corner = getCornerAt(x, y);
canvas.style.cursor = corner ? 'move' : 'crosshair';
});
canvas.addEventListener('mouseup', () => {
if (creatingBox && createStart && dragStart) {
const x1 = Math.min(createStart.x, dragStart.x);
const y1 = Math.min(createStart.y, dragStart.y);
const x2 = Math.max(createStart.x, dragStart.x);
const y2 = Math.max(createStart.y, dragStart.y);
if (x2 - x1 > 10 && y2 - y1 > 10) {
boxes.push({
bbox: [x1, y1, x2, y2],
label: 'knot',
confidence: 1.0,
source: 'manual',
});
syncHidden();
redraw();
}
}
isDragging = false;
creatingBox = false;
selectedCorner = null;
selectedBoxIndex = -1;
createStart = null;
dragStart = null;
canvas.style.cursor = 'crosshair';
});
})();
"""
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()
@ -426,244 +633,25 @@ class AnnotationApp:
img.save(buffer, format="PNG", optimize=True)
img_base64 = base64.b64encode(buffer.getvalue()).decode()
# Generate HTML with canvas and JavaScript
# Build HTML without <script> tags (Gradio sanitizes them). The interactive
# canvas logic runs via CANVAS_JS_ON_LOAD.
boxes_json = json.dumps(boxes)
html = f"""
<div style="position: relative; display: inline-block; border: 1px solid #ccc; padding: 5px;">
<div>Canvas Size: {display_width}x{display_height}</div>
<canvas id="annotation-canvas" width="{display_width}" height="{display_height}"
style="border: 1px solid #ccc; cursor: crosshair; max-width: 100%; height: auto; background-color: #f0f0f0;"></canvas>
<div id="canvas-info" style="margin-top: 5px; font-size: 12px; color: #666;">
Canvas ready - testing image loading
boxes_escaped = html.escape(boxes_json)
html_out = f"""
<div style="display: inline-block; border: 1px solid #ccc; padding: 5px;">
<div style="font-size: 12px; color: #666; margin-bottom: 4px;">Canvas Size: {display_width}x{display_height}</div>
<textarea id="annotation-initial-boxes" style="display:none;">{boxes_escaped}</textarea>
<div style="position: relative; width: {display_width}px; height: {display_height}px;">
<img id="annotation-img" src="data:image/png;base64,{img_base64}"
style="position:absolute; left:0; top:0; width:{display_width}px; height:{display_height}px;" />
<canvas id="annotation-canvas" width="{display_width}" height="{display_height}"
style="position:absolute; left:0; top:0; width:{display_width}px; height:{display_height}px; cursor: crosshair; background: transparent;"></canvas>
</div>
</div>
<script>
(function() {{
const canvas = document.getElementById('annotation-canvas');
const ctx = canvas.getContext('2d');
const img = new Image();
img.src = 'data:image/png;base64,{img_base64}';
let boxes = {boxes_json};
let isDragging = false;
let dragStart = null;
let selectedCorner = null;
let selectedBoxIndex = -1;
let creatingBox = false;
let createStart = null;
// Handle size for high DPI displays
const devicePixelRatio = window.devicePixelRatio || 1;
const displayWidth = {display_width};
const displayHeight = {display_height};
// Temporarily disable high DPI scaling to fix image display
// canvas.width = displayWidth * devicePixelRatio;
// canvas.height = displayHeight * devicePixelRatio;
// canvas.style.width = displayWidth + 'px';
// canvas.style.height = displayHeight + 'px';
// ctx.scale(devicePixelRatio, devicePixelRatio);
img.onload = function() {{
console.log('Image loaded successfully, size:', img.width, 'x', img.height);
redraw();
}};
img.onerror = function() {{
console.error('Failed to load image, drawing test pattern');
// Draw a test pattern if image fails to load
ctx.fillStyle = '#ffcccc';
ctx.fillRect(0, 0, displayWidth, displayHeight);
ctx.fillStyle = 'black';
ctx.font = '20px Arial';
ctx.fillText('Image failed to load', 10, 30);
ctx.fillText('Canvas size: ' + displayWidth + 'x' + displayHeight, 10, 60);
}};
// Force redraw after a short delay
setTimeout(function() {{
console.log('Timeout check - image complete:', img.complete, 'natural size:', img.naturalWidth, 'x', img.naturalHeight);
if (!img.complete || img.naturalWidth === 0) {{
console.log('Image not loaded, drawing test pattern');
ctx.fillStyle = '#ccccff';
ctx.fillRect(0, 0, displayWidth, displayHeight);
ctx.fillStyle = 'black';
ctx.font = '16px Arial';
ctx.fillText('Waiting for image...', 10, 30);
ctx.fillText('Canvas: ' + displayWidth + 'x' + displayHeight, 10, 50);
}}
}}, 500);
function redraw() {{
ctx.clearRect(0, 0, displayWidth, displayHeight);
ctx.drawImage(img, 0, 0, displayWidth, displayHeight);
// Draw boxes
boxes.forEach((box, index) => {{
const [x1, y1, x2, y2] = box.bbox;
const label = box.label || 'knot';
const conf = box.confidence || 1.0;
// Draw box
ctx.strokeStyle = 'red';
ctx.lineWidth = 3;
ctx.strokeRect(x1, y1, x2 - x1, y2 - y1);
// Draw corner handles
const handleSize = 6;
ctx.fillStyle = 'red';
ctx.fillRect(x1 - handleSize, y1 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillRect(x2 - handleSize, y1 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillRect(x1 - handleSize, y2 - handleSize, handleSize * 2, handleSize * 2);
ctx.fillRect(x2 - handleSize, y2 - handleSize, handleSize * 2, handleSize * 2);
// Draw label
ctx.fillStyle = 'red';
ctx.font = '16px Arial';
const text = conf < 1.0 ? `${{label}} ${{conf.toFixed(2)}}` : label;
ctx.fillText(text, x1, y1 - 5);
}});
// Draw creation box if creating
if (creatingBox && createStart && dragStart) {{
ctx.strokeStyle = 'blue';
ctx.lineWidth = 2;
ctx.setLineDash([5, 5]);
const x = Math.min(createStart.x, dragStart.x);
const y = Math.min(createStart.y, dragStart.y);
const w = Math.abs(createStart.x - dragStart.x);
const h = Math.abs(createStart.y - dragStart.y);
ctx.strokeRect(x, y, w, h);
ctx.setLineDash([]);
}}
}}
function getCornerAt(x, y) {{
const handleSize = 6;
for (let i = 0; i < boxes.length; i++) {{
const [x1, y1, x2, y2] = boxes[i].bbox;
const corners = [
{{x: x1, y: y1, type: 'top-left'}},
{{x: x2, y: y1, type: 'top-right'}},
{{x: x1, y: y2, type: 'bottom-left'}},
{{x: x2, y: y2, type: 'bottom-right'}}
];
for (let corner of corners) {{
if (x >= corner.x - handleSize && x <= corner.x + handleSize &&
y >= corner.y - handleSize && y <= corner.y + handleSize) {{
return {{boxIndex: i, corner: corner.type, pos: corner}};
}}
}}
}}
return null;
}}
canvas.addEventListener('mousedown', function(e) {{
const rect = canvas.getBoundingClientRect();
const x = (e.clientX - rect.left) * (displayWidth / rect.width);
const y = (e.clientY - rect.top) * (displayHeight / rect.height);
selectedCorner = getCornerAt(x, y);
if (selectedCorner) {{
isDragging = true;
selectedBoxIndex = selectedCorner.boxIndex;
canvas.style.cursor = 'move';
}} else {{
// Start creating new box
creatingBox = true;
createStart = {{x: x, y: y}};
dragStart = {{x: x, y: y}};
canvas.style.cursor = 'crosshair';
}}
}});
canvas.addEventListener('mousemove', function(e) {{
const rect = canvas.getBoundingClientRect();
const x = (e.clientX - rect.left) * (displayWidth / rect.width);
const y = (e.clientY - rect.top) * (displayHeight / rect.height);
if (isDragging && selectedCorner) {{
const box = boxes[selectedBoxIndex];
if (selectedCorner.corner === 'top-left') {{
box.bbox[0] = Math.min(x, box.bbox[2] - 10); // x1
box.bbox[1] = Math.min(y, box.bbox[3] - 10); // y1
}} else if (selectedCorner.corner === 'top-right') {{
box.bbox[2] = Math.max(x, box.bbox[0] + 10); // x2
box.bbox[1] = Math.min(y, box.bbox[3] - 10); // y1
}} else if (selectedCorner.corner === 'bottom-left') {{
box.bbox[0] = Math.min(x, box.bbox[2] - 10); // x1
box.bbox[3] = Math.max(y, box.bbox[1] + 10); // y2
}} else if (selectedCorner.corner === 'bottom-right') {{
box.bbox[2] = Math.max(x, box.bbox[0] + 10); // x2
box.bbox[3] = Math.max(y, box.bbox[1] + 10); // y2
}}
document.getElementById('canvas-boxes-data').value = JSON.stringify(boxes);
redraw();
}} else if (creatingBox && createStart) {{
dragStart = {{x: x, y: y}};
redraw();
}} else {{
// Update cursor
const corner = getCornerAt(x, y);
if (corner) {{
canvas.style.cursor = 'move';
}} else {{
canvas.style.cursor = 'crosshair';
}}
}}
}});
canvas.addEventListener('mouseup', function(e) {{
if (creatingBox && createStart && dragStart) {{
const x1 = Math.min(createStart.x, dragStart.x);
const y1 = Math.min(createStart.y, dragStart.y);
const x2 = Math.max(createStart.x, dragStart.x);
const y2 = Math.max(createStart.y, dragStart.y);
// Only add if box has minimum size
if (x2 - x1 > 10 && y2 - y1 > 10) {{
boxes.push({{
bbox: [x1, y1, x2, y2],
label: 'knot',
confidence: 1.0,
source: 'manual'
}});
document.getElementById('canvas-boxes-data').value = JSON.stringify(boxes);
redraw();
// Trigger save (this will be handled by the parent)
if (window.saveAnnotations) {{
window.saveAnnotations(boxes);
}}
}}
}}
isDragging = false;
creatingBox = false;
selectedCorner = null;
selectedBoxIndex = -1;
createStart = null;
dragStart = null;
canvas.style.cursor = 'crosshair';
}});
// Make boxes available globally for saving
window.getCurrentBoxes = function() {{
return JSON.stringify(boxes);
}};
// Function to update boxes from external call
window.updateBoxes = function(newBoxes) {{
boxes = newBoxes;
redraw();
}};
}})();
</script>
"""
return html
return html_out
def _format_boxes_text(self, boxes: list[dict]) -> str:
"""Format boxes for display."""
@ -1282,7 +1270,11 @@ def create_ui(app: AnnotationApp) -> gr.Blocks:
with gr.Row():
with gr.Column(scale=3):
image_display = gr.HTML(label="Current Image", value="<div style='width: 800px; height: 400px; border: 1px solid #ccc; display: flex; align-items: center; justify-content: center; color: #666;'>Load images from Settings to start annotating</div>")
image_display = gr.HTML(
label="Current Image",
value="<div style='width: 800px; height: 400px; border: 1px solid #ccc; display: flex; align-items: center; justify-content: center; color: #666;'>Load images from Settings to start annotating</div>",
js_on_load=CANVAS_JS_ON_LOAD,
)
with gr.Row():
prev_btn = gr.Button("⬅️ Previous")
@ -1291,7 +1283,7 @@ def create_ui(app: AnnotationApp) -> gr.Blocks:
save_canvas_btn = gr.Button("💾 Save Canvas Changes")
# Hidden textbox to store canvas boxes data
canvas_boxes_data = gr.Textbox(visible=False)
canvas_boxes_data = gr.Textbox(visible=False, elem_id="canvas-boxes-data")
with gr.Row():
threshold_slider = gr.Slider(0.1, 0.9, DEFAULT_DETECTION_THRESHOLD, label="Detection Threshold")
@ -1614,6 +1606,12 @@ 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 to run the GUI on"
)
args = parser.parse_args()
# Validate paths if provided
@ -1651,7 +1649,7 @@ def main():
demo.launch(
server_name="0.0.0.0",
server_port=7860,
server_port=args.port,
share=False
)