mx/AareLC_train
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity
Files
391c357c84a46f707d7d886a260e44d32c04f9bf
AareLC_train/src/gui/ml_gui.py
T
duan_j 391c357c84 migration and splitting AareLC
2026-04-14 16:07:31 +02:00

2523 lines
109 KiB
Python
Raw Blame History

"""AareLC ML Studio - Main GUI Application."""
import os
import time
import json
import tkinter as tk
import urllib.parse
from tkinter import ttk
from tkinter import filedialog, messagebox, simpledialog
import cv2
import requests
from collections import deque
import numpy as np
from pathlib import Path
from datetime import datetime
from PIL import Image, ImageTk
# Import modular components
from src.gui.canvas_panel import MainCanvas
from src.gui.control_panel import ControlPanel
from src.gui.class_bank_mixin import ClassBankMixin
from src.gui.dataset_prep_mixin import DatasetPrepMixin
from src.gui.canvas_interaction_mixin import CanvasInteractionMixin
from src.gui.review_mixin import ReviewMixin
from src.core.zmq_client import ZMQStreamClient
from src.core.inference_client import InferenceClient
from src.core.db_client import DatabaseClient
from src.core.image_processor import ImageProcessor
AAREDB_SHARED_PASSWORD = os.getenv('AAREDB_SHARED_PASSWORD')
ZMQ_CONTROL_TOKEN = os.getenv('AARELC_ZMQ_CONTROL_TOKEN')
class InferenceGUI(ClassBankMixin, DatasetPrepMixin, CanvasInteractionMixin, ReviewMixin):
"""Main GUI application for ML inference and annotation."""
def __init__(self, root):
self.root = root
self.root.title("AareLC ML Studio")
self.root.geometry("1400x900")
# --- Clients ---
self.zmq_client = None
self.inference_client = None
self.db_client = None
self.image_processor = ImageProcessor()
# --- Settings ---
self.pref_inference_url = tk.StringVar(value="http://mx-ml.psi.ch:8002")
self.pref_db_url = tk.StringVar(value="https://mx-db-01.psi.ch")
self.pref_download_url = tk.StringVar(value="http://localhost:8002/backend")
self.pref_zmq_control_url = tk.StringVar(value="http://localhost:8090")
self.pref_zmq_control_token = tk.StringVar(value=ZMQ_CONTROL_TOKEN)
self.pref_shared_pw = tk.StringVar(value=AAREDB_SHARED_PASSWORD)
self.pref_username = tk.StringVar(value=os.getenv('USER', 'anonymous'))
self.zmq_det_url = tk.StringVar(value="tcp://localhost:9091")
self.focus_enabled_var = tk.BooleanVar(value=True)
self.pref_zmq_stream_options = [
"tcp://localhost:9091",
"tcp://x10sa-spark-01:9091",
"tcp://sls-gpu-003:9089",
]
self.pref_zmq_control_options = [
"http://localhost:8090",
"http://x10sa-spark-01:8090",
"http://sls-gpu-003:8090",
]
self.pref_download_url_options = [
"http://localhost:8002/backend",
"http://mx-db-01:8002/backend",
]
# --- Sync filters ---
self.sync_since = tk.StringVar(value="")
self.sync_until = tk.StringVar(value="")
self.sync_limit = tk.IntVar(value=5000)
self.sync_include_unannotated = tk.BooleanVar(value=False)
self.sync_dry_run = tk.BooleanVar(value=True)
self.review_mode = False
self.review_items = []
self.review_index = -1
self.current_review_item = None
self.review_format = None
self.review_coco_path = None
self.review_coco_data = None
# --- State ---
self.current_raw_frame = None
self.current_image_id = None # Track current image ID from database
self.current_annotation_id = None # Track annotation ID for this image
self.last_results = []
self._last_overlap_pick = None
self.selected_island_idx = 0
self.selected_island_kind = "outer"
self.history = deque(maxlen=20) # NEW: Store last 20 states
self.focus_history = deque(maxlen=500)
self.stream_out_timestamps = deque(maxlen=240)
self.stream_fps_in = 0.0
self.stream_fps_out = 0.0
self.stream_fps_pred = 0.0
self._last_server_frame_id = None
self._last_server_frame_ts = None
# UI render cap for live stream; 0 means uncapped while still dropping stale frames.
# Override with env var AARELC_STREAM_UI_MAX_FPS (e.g. 8, 15, 0).
try:
self.stream_ui_max_fps = float(os.getenv("AARELC_STREAM_UI_MAX_FPS", "0"))
except ValueError:
self.stream_ui_max_fps = 0.0
if self.stream_ui_max_fps > 0:
print(f"INFO: Stream UI FPS cap enabled at {self.stream_ui_max_fps:.1f} FPS", flush=True)
else:
print("INFO: Stream UI FPS cap disabled (uncapped)", flush=True)
self._last_ui_render_ts = 0.0
self._pending_stream_frame = None
self._pending_stream_detections = None
self._pending_stream_metadata = None
self._last_stream_metadata = {}
self._stream_render_scheduled = False
# Status Bar Variables
self.status_var = tk.StringVar(value="Ready")
# Setup status truncation early to prevent X11 errors
self._setup_status_truncation()
self.local_output_dir = Path("data/custom_annotations")
self.current_local_image_path = None
self.class_set_source_path = None
self.element_bank_path = self.local_output_dir / "element_bank.json"
self.element_bank = []
self.annotation_mode = "yolo_single_polygon"
self.annotation_mode_labels = {
"yolo_single_polygon": "YOLO single polygon",
"coco_multi_polygon": "COCO multi-polygon",
}
# YOLO segmentation classes
self.class_params = {
"loop_all": {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01},
"pin": {"low": 100, "high": 200, "clahe": 2.0, "morph": 5, "eps": 0.002},
"crystal": {"low": 50, "high": 150, "clahe": 3.0, "morph": 11, "eps": 0.01},
"loop_face": {"low": 30, "high": 100, "clahe": 4.0, "morph": 3, "eps": 0.005},
"ice": {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01},
"needle": {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01}
}
# Color scheme mapping (BGR for OpenCV) for YOLO classes
self.class_colors = {
"loop_all": (0, 255, 0), # Blue
"pin": (0, 0, 255), # Green
"crystal": (255, 0, 0), # Red
"loop_face": (0, 255, 255), # Yellow
"ice": (255, 255, 255), # White/Cyan
"needle": (128, 0, 255)
}
# Map class names to YOLO class IDs (0-indexed for YOLO format)
self.class_to_id = {
"loop_all": 0,
"pin": 1,
"crystal": 2,
"loop_face": 3,
"ice": 4,
"needle": 5,
}
self.id_to_class = {v: k for k, v in self.class_to_id.items()}
self._class_color_palette = [
(0, 255, 0),
(0, 0, 255),
(255, 0, 0),
(0, 255, 255),
(255, 255, 255),
(128, 0, 255),
(255, 165, 0),
(255, 0, 255),
(0, 128, 255),
(128, 255, 0),
]
self._visibility_trace_bound_classes = set()
# Tkinter variables for control panel
self.tk_vars = {
'active_class': tk.StringVar(value=self._default_class_name()),
'active_tool': tk.StringVar(value="select"), # select, brush, sam, or measure
'canny_high': tk.IntVar(value=150),
'morph_kernel': tk.IntVar(value=7),
'poly_epsilon': tk.DoubleVar(value=0.01),
'show_debug_edges': tk.BooleanVar(value=False),
'plot_window_size': tk.IntVar(value=100),
'annotation_line_width': tk.IntVar(value=2),
'show_boxes': tk.BooleanVar(value=True),
'show_segments': tk.BooleanVar(value=True),
'brush_mode': tk.BooleanVar(value=False),
'brush_type': tk.StringVar(value='background'),
'brush_size': tk.IntVar(value=15),
'segmentation_method': tk.StringVar(value='smart') # Default to smart!
}
# Brush strokes storage: detection_idx -> {'bg': [], 'fg': []}
self.brush_strokes = {}
self.analysis_state = {
'show_grid': False,
'show_crosshair': True,
'cursor': (0, 0),
'line_start': None,
'line_end': None,
'line_dragging': False,
'roi_start': None,
'roi_end': None,
'roi_dragging': False,
}
self.depth_backend_var = tk.StringVar(value=os.getenv("AARELC_DEPTH_BACKEND", "depthpro"))
self.depth_max_side_var = tk.IntVar(value=int(os.getenv("AARELC_DEPTH_MAX_SIDE", "0")))
self.depth_preview_window = None
self.depth_preview_label = None
self.depth_preview_photo = None
# Initialize clients
self._init_clients()
# Setup UI
self._setup_menus()
self._setup_ui()
self._load_default_element_bank_if_exists()
# Status bar
self.status_bar = tk.Label(
self.root,
textvariable=self.status_var,
bd=1,
relief=tk.SUNKEN,
anchor=tk.W,
font="fixed"
)
self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)
# Setup parameter traces
self._setup_traces()
print(f"AAREDB_SHARED_PASSWORD={AAREDB_SHARED_PASSWORD}")
print(f"AARELC_ZMQ_CONTROL_TOKEN={ZMQ_CONTROL_TOKEN}")
def _init_clients(self):
"""Initialize client objects."""
self.inference_client = InferenceClient(self.pref_inference_url.get())
self.db_client = DatabaseClient(
self.pref_db_url.get(),
self.pref_shared_pw.get(),
download_base=self.pref_download_url.get()
)
def _setup_menus(self):
"""Setup menu bar."""
menubar = tk.Menu(self.root)
# File menu
file_menu = tk.Menu(menubar, tearoff=0)
file_menu.add_command(label="Load Image from Disk...", command=self.load_local_image)
file_menu.add_command(label="Fetch Next from DB", command=self.fetch_db_image)
file_menu.add_separator()
file_menu.add_command(label="Set Local Save Folder...", command=self.choose_local_output_dir)
file_menu.add_separator()
file_menu.add_command(label="Add Class...", command=self.add_single_class)
file_menu.add_command(label="Create New Class Set...", command=self.create_new_class_set)
file_menu.add_command(label="Load Class Set from Disk...", command=self.load_class_set_from_disk)
file_menu.add_command(label="Save Class Set As...", command=self.save_class_set_as)
file_menu.add_separator()
file_menu.add_command(label="Exit", command=self.root.quit)
menubar.add_cascade(label="File", menu=file_menu)
# Preferences menu
pref_menu = tk.Menu(menubar, tearoff=0)
pref_menu.add_command(label="Server Settings...", command=self.show_preferences)
pref_menu.add_command(label="ZMQ Inference Control...", command=self.show_zmq_inference_control_dialog)
menubar.add_cascade(label="Preferences", menu=pref_menu)
# Sync menu
sync_menu = tk.Menu(menubar, tearoff=0)
sync_menu.add_command(
label="Retrieve Images + Labels Locally...",
command=self.show_sync_dialog,
)
sync_menu.add_command(
label="Generate YOLO Segmentation Prep Script...",
command=self.show_dataset_prep_script_dialog,
)
menubar.add_cascade(label="Sync", menu=sync_menu)
# Review menu
review_menu = tk.Menu(menubar, tearoff=0)
review_menu.add_command(
label="Open Review Folder...",
command=self.open_review_folder,
)
review_menu.add_command(label="Previous Review Image (P)", command=self.prev_review_image)
review_menu.add_command(label="Next Review Image (N)", command=self.next_review_image)
review_menu.add_separator()
review_menu.add_command(label="Save Current Review Annotation", command=self.save_review_annotation)
review_menu.add_command(label="Exit Review Mode", command=self.exit_review_mode)
menubar.add_cascade(label="Review", menu=review_menu)
# Elements menu
elements_menu = tk.Menu(menubar, tearoff=0)
elements_menu.add_command(label="Add Selected to Bank...", command=self.add_selected_to_element_bank)
elements_menu.add_command(label="Place Element from Bank...", command=self.place_element_from_bank)
elements_menu.add_separator()
elements_menu.add_command(label="Load Element Bank...", command=self.load_element_bank_from_disk)
elements_menu.add_command(label="Save Element Bank As...", command=self.save_element_bank_as)
menubar.add_cascade(label="Elements", menu=elements_menu)
self.root.config(menu=menubar)
def _setup_ui(self):
"""Setup main UI layout."""
# 1. Top toolbars (two rows to avoid clipping on smaller screens)
self.top_bar = tk.Frame(self.root, bg="#eee", pady=5, relief=tk.RAISED, bd=1)
self.top_bar.pack(side=tk.TOP, fill=tk.X)
self.top_bar_aux = tk.Frame(self.root, bg="#f3f3f3", pady=4, relief=tk.RAISED, bd=1)
self.top_bar_aux.pack(side=tk.TOP, fill=tk.X)
# Removed ZMQ Entry from here
# NEW: Tool Palette Frame (Left side of toolbar)
self.tool_frame = tk.Frame(self.top_bar, bg="#ddd", padx=5)
self.tool_frame.pack(side=tk.LEFT, padx=5)
self.select_tool_btn = tk.Radiobutton(
self.tool_frame, text="Select", variable=self.tk_vars['active_tool'],
value="select", indicatoron=0, width=8, padx=5, pady=2
)
self.select_tool_btn.pack(side=tk.LEFT)
self.brush_tool_btn = tk.Radiobutton(
self.tool_frame, text="Brush", variable=self.tk_vars['active_tool'],
value="brush", indicatoron=0, width=8, padx=5, pady=2
)
self.brush_tool_btn.pack(side=tk.LEFT)
self.sam_tool_btn = tk.Radiobutton(
self.tool_frame, text="SAM", variable=self.tk_vars['active_tool'],
value="sam", indicatoron=0, width=8, padx=5, pady=2
)
self.sam_tool_btn.pack(side=tk.LEFT)
self.measure_tool_btn = tk.Radiobutton(
self.tool_frame, text="Measure", variable=self.tk_vars['active_tool'],
value="measure", indicatoron=0, width=10, padx=5, pady=2
)
self.measure_tool_btn.pack(side=tk.LEFT)
tk.Label(self.top_bar, text="|", bg="#eee").pack(side=tk.LEFT, padx=5)
# Connect button
self.stream_btn = tk.Button(self.top_bar, text="Connect Live", command=self.toggle_stream, width=15)
self.stream_btn.pack(side=tk.LEFT, padx=5)
self.focus_btn = tk.Button(
self.top_bar,
text="Focus ON",
command=self.toggle_focus_stream,
width=12,
bg="#17a2b8",
fg="white"
)
self.focus_btn.pack(side=tk.LEFT, padx=5)
# Database Fetch button (Next Random Image)
self.db_next_btn = tk.Button(self.top_bar, text="Next Image", command=self.fetch_db_image, bg="#28a745",
fg="white")
self.db_next_btn.pack(side=tk.LEFT, padx=5)
# Skip Image button
self.skip_btn = tk.Button(self.top_bar, text="Skip", command=self.skip_current_image, bg="#ffc107",
fg="black")
self.skip_btn.pack(side=tk.LEFT, padx=2)
# Predict button
self.predict_btn = tk.Button(self.top_bar, text="Predict (/predict)", command=self.request_inference,
bg="#007bff", fg="white")
self.predict_btn.pack(side=tk.LEFT, padx=20)
self.depth_btn = tk.Button(
self.top_bar,
text="Depth Map",
command=self.request_depth_map,
bg="#5a3d00",
fg="white"
)
self.depth_btn.pack(side=tk.LEFT, padx=5)
self.sam_boxes_btn = tk.Button(
self.top_bar_aux,
text="SAM from Boxes",
command=self.refine_boxes_with_sam,
bg="#17a2b8",
fg="white"
)
self.sam_boxes_btn.pack(side=tk.LEFT, padx=5)
tk.Label(self.top_bar_aux, text="Depth:", bg="#f3f3f3").pack(side=tk.LEFT, padx=(8, 2))
self.depth_backend_combo = ttk.Combobox(
self.top_bar_aux,
textvariable=self.depth_backend_var,
values=[
"depthpro",
"depth_anything_v1",
"depth_anything_v1_trt",
"depth_anything_v2",
"depth_anything_v2_tiny",
"depth_anything_v2_trt",
"midas_v21_small",
],
width=20,
state="readonly"
)
self.depth_backend_combo.pack(side=tk.LEFT, padx=2)
tk.Label(self.top_bar_aux, text="Max side:", bg="#f3f3f3").pack(side=tk.LEFT, padx=(6, 2))
self.depth_max_side_combo = ttk.Combobox(
self.top_bar_aux,
textvariable=self.depth_max_side_var,
values=[0, 384, 512, 640, 768, 1024],
width=6,
state="readonly"
)
self.depth_max_side_combo.pack(side=tk.LEFT, padx=2)
# NEW: Zoom Slider in Toolbar
tk.Label(self.top_bar_aux, text="Zoom:", bg="#f3f3f3").pack(side=tk.LEFT, padx=(10, 2))
self.zoom_var = tk.DoubleVar(value=1.0)
self.zoom_slider = tk.Scale(
self.top_bar_aux, from_=0.1, to=5.0, resolution=0.1,
orient=tk.HORIZONTAL, variable=self.zoom_var,
showvalue=False, width=10, length=120, bg="#f3f3f3",
command=self._on_zoom_slider_change
)
self.zoom_slider.pack(side=tk.LEFT, padx=5)
# NEW: SAM refinement toggle
self.use_sam_refine = tk.BooleanVar(value=False)
self.sam_refine_check = tk.Checkbutton(
self.top_bar_aux, text="Use SAM Refinement", variable=self.use_sam_refine,
bg="#f3f3f3", font=('Arial', 9)
)
self.sam_refine_check.pack(side=tk.LEFT, padx=5)
self.show_boxes_check = tk.Checkbutton(
self.top_bar_aux,
text="Show Boxes",
variable=self.tk_vars['show_boxes'],
bg="#f3f3f3",
font=('Arial', 9)
)
self.show_boxes_check.pack(side=tk.LEFT, padx=4)
self.show_segments_check = tk.Checkbutton(
self.top_bar_aux,
text="Show Segments",
variable=self.tk_vars['show_segments'],
bg="#f3f3f3",
font=('Arial', 9)
)
self.show_segments_check.pack(side=tk.LEFT, padx=4)
# 2. Main content area
self.content_frame = tk.Frame(self.root)
self.content_frame.pack(fill=tk.BOTH, expand=True)
# 3. Control panel (right side)
self.ctrl_panel = ControlPanel(
self.content_frame,
self.class_params,
self.tk_vars,
callbacks={
'on_train_click': self.send_to_training,
'on_clear_brush': self.clear_brush_strokes,
'on_regenerate_edge': self.regenerate_selected_edge
}
)
self.ctrl_panel.pack(side=tk.RIGHT, fill=tk.Y)
# Initialize class visibility checkboxes
self.ctrl_panel.update_class_visibility(list(self.class_params.keys()), self.class_colors)
self.ctrl_panel.update_class_selector(list(self.class_params.keys()))
self._bind_class_visibility_traces()
# Add trace to redraw when brush mode changes
self.tk_vars['brush_mode'].trace_add("write", lambda *args: self.redraw_current_frame())
# 4. Canvas (center)
self.canvas = MainCanvas(self.content_frame)
# FIX: Use _on_canvas_interaction, which handles click/drag arguments
self.canvas.set_redraw_callback(self._on_canvas_interaction)
self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
self.add_box_btn = tk.Button(
self.top_bar,
text="Add Box",
command=self.add_new_box,
bg="#6c757d",
fg="white"
)
self.add_box_btn.pack(side=tk.LEFT, padx=5)
self.add_island_btn = tk.Button(
self.top_bar_aux,
text="Island",
command=self.add_island_to_selected,
bg="#6c757d",
fg="white"
)
self.add_island_btn.pack(side=tk.LEFT, padx=2)
self.add_hole_btn = tk.Button(
self.top_bar_aux,
text="Add Hole",
command=self.add_hole_to_selected,
bg="#6c757d",
fg="white"
)
self.add_hole_btn.pack(side=tk.LEFT, padx=2)
self.next_island_btn = tk.Button(
self.top_bar_aux,
text="Cycle Island",
command=self.cycle_selected_island,
bg="#f8f9fa"
)
self.next_island_btn.pack(side=tk.LEFT, padx=2)
self.remove_island_btn = tk.Button(
self.top_bar_aux,
text="Remove Island",
command=self.remove_selected_island,
bg="#f8f9fa"
)
self.remove_island_btn.pack(side=tk.LEFT, padx=2)
# NEW: Undo Button
self.undo_btn = tk.Button(
self.top_bar,
text="Undo",
command=self.undo_action,
bg="#f8f9fa"
)
self.undo_btn.pack(side=tk.LEFT, padx=5)
def _on_zoom_slider_change(self, val):
"""Sync slider value to canvas zoom level."""
self.canvas.zoom_level = float(val)
self.redraw_current_frame()
def _setup_status_truncation(self):
"""Setup status message truncation to prevent X11 BadLength errors."""
def truncate_status(*args):
val = self.status_var.get()
if len(val) > 100:
truncated = val[:97] + "..."
if val != truncated:
self.status_var.set(truncated)
self.status_var.trace_add("write", truncate_status)
def _setup_traces(self):
"""Setup variable traces for parameter changes."""
self._after_id = None # NEW: Track the debouncing timer
for var_name in ['canny_high', 'morph_kernel', 'poly_epsilon']:
self.tk_vars[var_name].trace_add("write", self.on_param_change)
self.tk_vars['active_class'].trace_add("write", self.on_class_switch)
self.tk_vars['active_tool'].trace_add("write", self._on_tool_switch)
self.tk_vars['show_debug_edges'].trace_add("write", lambda *args: self.redraw_current_frame())
self.tk_vars['plot_window_size'].trace_add("write", lambda *args: self.redraw_focus_plot())
self.tk_vars['annotation_line_width'].trace_add("write", lambda *args: self.redraw_current_frame())
self.tk_vars['show_boxes'].trace_add("write", lambda *args: self.redraw_current_frame())
self.tk_vars['show_segments'].trace_add("write", lambda *args: self.redraw_current_frame())
self.root.bind("<a>", lambda e: self.add_new_box())
self.root.bind("<Delete>", self.delete_selected)
self.root.bind("<Control-z>", lambda e: self.undo_action()) # NEW: Ctrl+Z shortcut
self.root.bind("<Command-z>", lambda e: self.undo_action()) # NEW: Cmd+Z for macOS
self.root.bind("<i>", lambda e: self.cycle_selected_island())
# Review navigation shortcuts
self.root.bind("<Right>", self._on_review_next_shortcut)
self.root.bind("<Left>", self._on_review_prev_shortcut)
self.root.bind("<n>", self._on_review_next_shortcut)
self.root.bind("<p>", self._on_review_prev_shortcut)
self.root.bind("<N>", self._on_review_next_shortcut)
self.root.bind("<P>", self._on_review_prev_shortcut)
self.root.bind("<g>", self._toggle_analysis_grid)
self.root.bind("<c>", self._toggle_analysis_crosshair)
self.root.bind("<x>", self._clear_measurements)
# Global safety: Truncate any status message to 100 chars to avoid X11 BadLength
# Moved to _setup_status_truncation for early setup
def _default_class_name(self):
"""Return first available class name."""
return next(iter(self.class_to_id.keys()), "class_0")
@staticmethod
def _default_class_params():
"""Default contour extraction params for newly created classes."""
return {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01}
def _is_coco_multi_mode(self):
"""Whether current class set uses COCO multi-polygon mode."""
return self.annotation_mode == "coco_multi_polygon"
def _get_annotation_mode_label(self):
return self.annotation_mode_labels.get(self.annotation_mode, self.annotation_mode)
def _prompt_annotation_mode(self):
"""Ask user which annotation mode should be used for a new class set."""
choice = messagebox.askyesnocancel(
"Annotation Mode",
(
"Enable COCO multi-polygon mode for this class set?\n\n"
"Yes: COCO multi-polygon (supports multi-island instances, local COCO save)\n"
"No: YOLO single-polygon (current/default behavior)\n"
"Cancel: abort"
),
parent=self.root,
)
if choice is None:
return None
return "coco_multi_polygon" if choice else "yolo_single_polygon"
def _det_get_polygons_relative(self, det):
"""Return list of relative polygons in (N,1,2) float32 format."""
polys = det.get("polygons")
out = []
if isinstance(polys, list):
for poly in polys:
arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
if len(arr) >= 3:
out.append(arr)
if out:
return out
poly = det.get("poly")
if poly is not None and len(poly) >= 3:
out.append(np.array(poly, dtype=np.float32).reshape(-1, 1, 2))
return out
def _det_get_holes_relative(self, det):
"""Return list of relative hole polygons in (N,1,2) float32 format."""
holes = det.get("holes")
out = []
if isinstance(holes, list):
for poly in holes:
arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
if len(arr) >= 3:
out.append(arr)
return out
def _det_has_polygon(self, det):
"""True if detection has at least one explicit polygon island."""
return len(self._det_get_polygons_relative(det)) > 0
def _set_det_single_polygon(self, det, poly_arr):
"""Set detection polygon in both legacy and multi-poly fields."""
arr = np.array(poly_arr, dtype=np.float32).reshape(-1, 1, 2)
det["poly"] = arr
det["polygons"] = [arr]
def _set_det_polygons(self, det, polygon_list):
"""Set detection polygons list, updating legacy primary poly too."""
cleaned = []
for poly in polygon_list:
arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
if len(arr) >= 3:
cleaned.append(arr)
if cleaned:
det["polygons"] = cleaned
det["poly"] = cleaned[0]
else:
det.pop("polygons", None)
det.pop("poly", None)
def _set_det_holes(self, det, hole_list):
"""Set detection hole polygons list."""
cleaned = []
for poly in hole_list:
arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
if len(arr) >= 3:
cleaned.append(arr)
if cleaned:
det["holes"] = cleaned
else:
det.pop("holes", None)
def _det_get_absolute_polygons(self, det):
"""Return polygon islands as absolute image coordinates."""
box = det.get("box", det)
bx1 = float(box.get("x1", 0.0))
by1 = float(box.get("y1", 0.0))
polygons_abs = []
for rel in self._det_get_polygons_relative(det):
pts = rel.reshape(-1, 2).astype(np.float32) + [bx1, by1]
if len(pts) >= 3:
polygons_abs.append(pts)
return polygons_abs
def _det_get_absolute_holes(self, det):
"""Return hole polygons as absolute image coordinates."""
box = det.get("box", det)
bx1 = float(box.get("x1", 0.0))
by1 = float(box.get("y1", 0.0))
holes_abs = []
for rel in self._det_get_holes_relative(det):
pts = rel.reshape(-1, 2).astype(np.float32) + [bx1, by1]
if len(pts) >= 3:
holes_abs.append(pts)
return holes_abs
def _get_active_polygon(self, det):
"""Get currently selected polygon island for a detection."""
rings = self._det_get_polygons_relative(det) if self.selected_island_kind == "outer" else self._det_get_holes_relative(det)
if not rings:
self.selected_island_kind = "outer"
rings = self._det_get_polygons_relative(det)
if not rings:
return None
idx = int(max(0, min(self.selected_island_idx, len(rings) - 1)))
self.selected_island_idx = idx
return rings[idx]
def add_island_to_selected(self):
"""Add a new polygon island inside selected detection box."""
idx = self.canvas.selected_obj_idx
if idx == -1 or idx >= len(self.last_results):
self.status_var.set("Select an instance first.")
return
det = self.last_results[idx]
box = det.get('box', det)
x1, y1, x2, y2 = float(box['x1']), float(box['y1']), float(box['x2']), float(box['y2'])
w = max(6.0, x2 - x1)
h = max(6.0, y2 - y1)
cx, cy = w * 0.5, h * 0.5
rw, rh = max(4.0, w * 0.2), max(4.0, h * 0.2)
poly = np.array(
[[[cx - rw, cy - rh]], [[cx + rw, cy - rh]], [[cx + rw, cy + rh]], [[cx - rw, cy + rh]]],
dtype=np.float32
)
polygons = self._det_get_polygons_relative(det)
polygons.append(poly)
self._save_history()
self._set_det_polygons(det, polygons)
self.selected_island_idx = len(polygons) - 1
self.canvas.selected_point_idx = -1
self.redraw_current_frame()
self.status_var.set(f"Added island {self.selected_island_idx + 1}/{len(polygons)}")
def add_hole_to_selected(self):
"""Add a new hole polygon inside selected detection box."""
idx = self.canvas.selected_obj_idx
if idx == -1 or idx >= len(self.last_results):
self.status_var.set("Select an instance first.")
return
det = self.last_results[idx]
box = det.get('box', det)
x1, y1, x2, y2 = float(box['x1']), float(box['y1']), float(box['x2']), float(box['y2'])
w = max(10.0, x2 - x1)
h = max(10.0, y2 - y1)
cx, cy = w * 0.5, h * 0.5
rw, rh = max(3.0, w * 0.12), max(3.0, h * 0.12)
hole = np.array(
[[[cx - rw, cy - rh]], [[cx + rw, cy - rh]], [[cx + rw, cy + rh]], [[cx - rw, cy + rh]]],
dtype=np.float32
)
holes = self._det_get_holes_relative(det)
holes.append(hole)
self._save_history()
self._set_det_holes(det, holes)
self.selected_island_kind = "hole"
self.selected_island_idx = len(holes) - 1
self.canvas.selected_point_idx = -1
self.redraw_current_frame()
self.status_var.set(f"Added hole {self.selected_island_idx + 1}/{len(holes)}")
def cycle_selected_island(self):
"""Cycle active island on selected detection."""
idx = self.canvas.selected_obj_idx
if idx == -1 or idx >= len(self.last_results):
self.status_var.set("Select an instance first.")
return
det = self.last_results[idx]
outers = self._det_get_polygons_relative(det)
holes = self._det_get_holes_relative(det)
if not outers and not holes:
self.status_var.set("No rings on selected instance.")
return
if self.selected_island_kind == "outer":
if outers:
self.selected_island_idx = (self.selected_island_idx + 1) % len(outers)
if self.selected_island_idx == 0 and holes:
self.selected_island_kind = "hole"
self.selected_island_idx = 0
else:
self.selected_island_kind = "hole"
self.selected_island_idx = (self.selected_island_idx + 1) % len(holes)
else:
if holes:
self.selected_island_idx = (self.selected_island_idx + 1) % len(holes)
if self.selected_island_idx == 0 and outers:
self.selected_island_kind = "outer"
self.selected_island_idx = 0
else:
self.selected_island_kind = "outer"
self.selected_island_idx = (self.selected_island_idx + 1) % len(outers)
self.canvas.selected_point_idx = -1
self.redraw_current_frame()
total = len(self._det_get_polygons_relative(det)) if self.selected_island_kind == "outer" else len(self._det_get_holes_relative(det))
self.status_var.set(f"Active {self.selected_island_kind} {self.selected_island_idx + 1}/{max(1, total)}")
def remove_selected_island(self):
"""Remove active ring from selected detection."""
idx = self.canvas.selected_obj_idx
if idx == -1 or idx >= len(self.last_results):
self.status_var.set("Select an instance first.")
return
det = self.last_results[idx]
if self.selected_island_kind == "hole":
holes = self._det_get_holes_relative(det)
if not holes:
self.status_var.set("No holes to remove.")
return
self._save_history()
del holes[self.selected_island_idx]
self.selected_island_idx = max(0, min(self.selected_island_idx, len(holes) - 1))
self._set_det_holes(det, holes)
if not holes:
self.selected_island_kind = "outer"
self.selected_island_idx = 0
else:
polygons = self._det_get_polygons_relative(det)
if not polygons:
self.status_var.set("No outer islands to remove.")
return
if len(polygons) == 1:
messagebox.showwarning("Remove Island", "Instance has a single outer island. Delete instance if needed.")
return
self._save_history()
del polygons[self.selected_island_idx]
self.selected_island_idx = max(0, min(self.selected_island_idx, len(polygons) - 1))
self._set_det_polygons(det, polygons)
self.canvas.selected_point_idx = -1
self.redraw_current_frame()
self.status_var.set(f"Removed ring. Active {self.selected_island_kind} {self.selected_island_idx + 1}")
# --- Event Handlers ---
def on_class_switch(self, *args):
"""Load stored parameters when the user switches classes in the dropdown."""
cls = self.tk_vars['active_class'].get()
p = self.class_params.get(cls, self._default_class_params())
# We use a flag to prevent on_param_change from being triggered
# while we are just updating the sliders to match the selected class
self._updating_ui = True
self.tk_vars['canny_high'].set(p["high"])
self.tk_vars['morph_kernel'].set(p["morph"])
self.tk_vars['poly_epsilon'].set(p["eps"])
self._updating_ui = False
self.redraw_current_frame()
def _on_tool_switch(self, *args):
"""Sync canvas interaction mode with selected tool."""
is_measure = self.tk_vars['active_tool'].get() == "measure"
self.canvas.suspend_pan = is_measure
if is_measure:
self.status_var.set("Measure mode: left-drag line, right-drag ROI, g/c/x toggles")
self.redraw_current_frame()
def on_param_change(self, *args):
"""Save slider values and CLEAR polygon ONLY for selected object."""
if getattr(self, '_updating_ui', False):
return
# Cancel previous scheduled redraw if user is still sliding
if self._after_id:
self.root.after_cancel(self._after_id)
def perform_update():
cls = self.tk_vars['active_class'].get()
self.class_params[cls].update({
"high": self.tk_vars['canny_high'].get(),
"morph": self.tk_vars['morph_kernel'].get(),
"eps": self.tk_vars['poly_epsilon'].get()
})
# ONLY clear polygon for the SELECTED object, not all objects of this class
if self.canvas.selected_obj_idx != -1 and self.canvas.selected_obj_idx < len(self.last_results):
det = self.last_results[self.canvas.selected_obj_idx]
label = det.get('name', det.get('label', self._default_class_name()))
# Only clear if the selected object matches the active class
if label == cls:
if 'poly' in det or 'polygons' in det:
det.pop('poly', None)
det.pop('polygons', None)
self.status_var.set(f"Parameters updated for selected {cls} - click Regenerate Edge")
else:
self.status_var.set(f"Note: Selected object is '{label}', not '{cls}'. Switch class to edit.")
else:
self.status_var.set(f"Parameters saved for class '{cls}' - will apply to new detections")
self.redraw_current_frame()
self._after_id = None
# Delay execution by 150ms to wait for slider to settle
self._after_id = self.root.after(150, perform_update)
def _save_history(self):
"""Save a deep copy of current detections to the history stack."""
import copy
# We store a copy of the list and the dictionaries inside
snapshot = copy.deepcopy(self.last_results)
self.history.append(snapshot)
def undo_action(self):
"""Restore the last saved state from history."""
if not self.history:
self.status_var.set("Nothing to undo")
return
self.last_results = self.history.pop()
self.canvas.selected_obj_idx = -1
self.selected_island_idx = 0
self.selected_island_kind = "outer"
self.redraw_current_frame()
self.status_var.set("Action undone")
def _run_sam_inference(self, x, y):
"""Send a point to the DGX SAM backend and create a new detection."""
if self.current_raw_frame is None: return
try:
self.root.config(cursor="watch")
self.status_var.set(f"SAM: Segmenting at [{int(x)}, {int(y)}]...")
self.root.update_idletasks()
result = self.inference_client.segment_anything(self.current_raw_frame, x, y)
poly_pts = result.get('polygon', [])
if not poly_pts:
self.status_var.set("SAM: No object found at that location.")
return
# Convert list of [x, y] to numpy format (N, 1, 2)
poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
# Calculate a bounding box for the polygon
x_coords = [p[0] for p in poly_pts]
y_coords = [p[1] for p in poly_pts]
x1, y1, x2, y2 = min(x_coords), min(y_coords), max(x_coords), max(y_coords)
# Polygons in our system are stored relative to box origin
poly_relative = poly_arr - [x1, y1]
new_det = {
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
"name": self.tk_vars['active_class'].get(),
"poly": poly_relative.astype(np.float32),
"polygons": [poly_relative.astype(np.float32)],
"conf": result.get('score', 1.0)
}
self._save_history()
self.last_results.append(new_det)
self.canvas.selected_obj_idx = len(self.last_results) - 1
self.selected_island_kind = "outer"
self.redraw_current_frame()
self.status_var.set(f"SAM: Object segmented successfully ({len(poly_pts)} points).")
except Exception as e:
messagebox.showerror("SAM Error", str(e))
self.status_var.set("SAM inference failed.")
finally:
self.root.config(cursor="")
# --- Image Loading ---
def load_local_image(self):
"""Load image from local file system."""
path = filedialog.askopenfilename(
filetypes=[("Images", "*.jpg *.jpeg *.png *.bmp")]
)
if path:
img = cv2.imread(path)
if img is not None:
self.exit_review_mode(silent=True)
self._stop_streaming()
self.current_raw_frame = img
self.current_image_id = None
self.current_annotation_id = None
self.current_local_image_path = Path(path)
self.last_results = []
self.selected_island_idx = 0
self.selected_island_kind = "outer"
self.redraw_current_frame()
self.status_var.set(f"Loaded: {path}")
def fetch_db_image(self):
"""Fetch image from database server."""
try:
self.exit_review_mode(silent=True)
self.root.config(cursor="watch")
# Fetch returns (img, metadata)
result = self.db_client.fetch_next_raw_image()
# Check if result is just an image or (image, metadata)
if isinstance(result, tuple):
img, metadata = result
self.current_image_id = metadata.get('image', {}).get('id')
self.current_annotation_id = metadata.get('annotation_id')
else:
# Backwards compatibility - just image
img = result
self.current_image_id = None
self.current_annotation_id = None
self._stop_streaming()
self.current_raw_frame = img
self.current_local_image_path = None
self.last_results = [] # Clear previous annotations
self.selected_island_idx = 0
self.selected_island_kind = "outer"
self.redraw_current_frame()
status_msg = f"Fetched image from database"
if self.current_image_id:
status_msg += f" (ID: {self.current_image_id})"
self.status_var.set(status_msg)
except Exception as e:
messagebox.showerror("DB Error", str(e))
self.status_var.set("Database fetch failed")
import traceback
traceback.print_exc()
finally:
self.root.config(cursor="")
# --- Inference ---
def request_inference(self):
"""Request /predict inference for the currently displayed frame snapshot."""
if self.current_raw_frame is None:
messagebox.showwarning("Warning", "No image loaded to predict.")
return
frame_for_predict = self.current_raw_frame.copy()
from_stream = isinstance(getattr(self, '_last_stream_metadata', None), dict) and bool(self._last_stream_metadata)
source_label = "live snapshot" if from_stream else "current image"
try:
self.root.config(cursor="watch")
self.root.update_idletasks()
self.status_var.set(f"Running /predict inference on {source_label}...")
self.root.update_idletasks()
self.last_results = self.inference_client.predict(frame_for_predict)
# Optional SAM refinement runs AFTER a fresh /predict call.
if self.use_sam_refine.get() and self.last_results:
refine_indices = [
i for i, det in enumerate(self.last_results)
if not self._det_has_polygon(det)
]
if refine_indices:
self.status_var.set(f"SAM: Auto-refining {len(refine_indices)} /predict detections...")
self.root.update_idletasks()
boxes_to_refine = []
for i in refine_indices:
b = self.last_results[i].get('box', self.last_results[i])
boxes_to_refine.append([int(b['x1']), int(b['y1']), int(b['x2']), int(b['y2'])])
sam_results = self.inference_client.segment_multi_box(frame_for_predict, boxes_to_refine)
for local_idx, res in enumerate(sam_results):
poly_pts = res.get('polygon', [])
if not poly_pts:
continue
det_idx = refine_indices[local_idx]
box = self.last_results[det_idx].get('box', self.last_results[det_idx])
poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
self._set_det_single_polygon(self.last_results[det_idx], poly_arr - [box['x1'], box['y1']])
self.last_results[det_idx].pop('suppress_auto_polygon', None)
self.redraw_current_frame()
self.root.update_idletasks()
self.status_var.set(f"Prediction complete via /predict/. Found {len(self.last_results)} objects.")
except Exception as e:
messagebox.showerror("Inference Error", str(e))
self.status_var.set("Inference failed")
finally:
self.root.config(cursor="")
def request_depth_map(self):
"""Request FastDepth map and show it in a dedicated preview window."""
if self.current_raw_frame is None:
messagebox.showwarning("Warning", "No image loaded for depth analysis.")
return
try:
self.root.config(cursor="watch")
backend = self.depth_backend_var.get().strip() or "depthpro"
depth_max_side = int(self.depth_max_side_var.get())
self.status_var.set(f"Running depth inference ({backend}, max_side={depth_max_side})...")
self.root.update_idletasks()
depth_img = self.inference_client.predict_depth(
self.current_raw_frame,
colorize=True,
depth_model=backend,
depth_max_side=depth_max_side
)
self._show_depth_preview(depth_img)
self.status_var.set(f"Depth map ready ({backend}, max_side={depth_max_side}).")
except Exception as e:
messagebox.showerror("Depth Error", str(e))
self.status_var.set("Depth inference failed")
finally:
self.root.config(cursor="")
def _show_depth_preview(self, depth_img):
"""Display the depth map in a non-modal window."""
if self.depth_preview_window is None or not self.depth_preview_window.winfo_exists():
self.depth_preview_window = tk.Toplevel(self.root)
self.depth_preview_window.title("FastDepth Preview")
self.depth_preview_window.geometry("900x700")
self.depth_preview_window.configure(bg="#111")
self.depth_preview_label = tk.Label(self.depth_preview_window, bg="#111")
self.depth_preview_label.pack(fill=tk.BOTH, expand=True, padx=8, pady=8)
h, w = depth_img.shape[:2]
max_w, max_h = 1200, 900
scale = min(max_w / max(1, w), max_h / max(1, h), 1.0)
if scale < 1.0:
depth_img = cv2.resize(depth_img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
if depth_img.ndim == 2:
rgb = cv2.cvtColor(depth_img, cv2.COLOR_GRAY2RGB)
else:
rgb = cv2.cvtColor(depth_img, cv2.COLOR_BGR2RGB)
pil = Image.fromarray(rgb)
self.depth_preview_photo = ImageTk.PhotoImage(pil)
self.depth_preview_label.configure(image=self.depth_preview_photo)
self.depth_preview_window.lift()
def refine_boxes_with_sam(self):
"""Send current rectangle boxes to SAM and update polygons."""
if self.current_raw_frame is None:
messagebox.showwarning("Warning", "No image loaded.")
return
if not self.last_results:
messagebox.showwarning("Warning", "No boxes available. Add or predict boxes first.")
return
try:
self.root.config(cursor="watch")
refine_indices = [
i for i, det in enumerate(self.last_results)
if not self._det_has_polygon(det)
]
if not refine_indices:
self.status_var.set("SAM: All boxes already have polygons; nothing to refine.")
return
self.status_var.set(f"SAM: Refining {len(refine_indices)} boxes...")
self.root.update_idletasks()
boxes_to_refine = []
for idx in refine_indices:
det = self.last_results[idx]
box = det.get('box', det)
boxes_to_refine.append([int(box['x1']), int(box['y1']), int(box['x2']), int(box['y2'])])
sam_results = self.inference_client.segment_multi_box(self.current_raw_frame, boxes_to_refine)
refined = 0
for local_idx, res in enumerate(sam_results):
poly_pts = res.get('polygon', [])
if not poly_pts:
continue
det_idx = refine_indices[local_idx]
box = self.last_results[det_idx].get('box', self.last_results[det_idx])
poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
self._set_det_single_polygon(self.last_results[det_idx], poly_arr - [box['x1'], box['y1']])
self.last_results[det_idx].pop('suppress_auto_polygon', None)
self.last_results[det_idx]['conf'] = res.get('score', self.last_results[det_idx].get('conf', 1.0))
refined += 1
self.redraw_current_frame()
self.status_var.set(f"SAM: Refined {refined}/{len(refine_indices)} boxes")
except Exception as e:
messagebox.showerror("SAM Error", str(e))
self.status_var.set("SAM box refinement failed")
finally:
self.root.config(cursor="")
# --- Streaming ---
def toggle_stream(self):
"""Toggle live streaming on/off."""
if self.zmq_client and self.zmq_client.is_streaming():
self._stop_streaming()
else:
self._start_streaming()
def _start_streaming(self):
"""Start ZMQ live streaming."""
self.stream_out_timestamps.clear()
self.stream_fps_in = 0.0
self.stream_fps_out = 0.0
self.stream_fps_pred = 0.0
self._last_server_frame_id = None
self._last_server_frame_ts = None
self.zmq_client = ZMQStreamClient(
self.zmq_det_url.get(),
on_frame_callback=self._on_zmq_frame,
on_focus_callback=self._on_focus_update
)
if not self.focus_enabled_var.get():
self.zmq_client.disable_focus()
self.zmq_client.start()
self.stream_btn.config(text="Stop Live", bg="red")
cuda_text = "ON" if getattr(self.zmq_client, 'cuda_enabled', False) else "OFF"
print(f"INFO: GUI stream CUDA focus: {cuda_text}", flush=True)
self.status_var.set(f"Streaming... CUDA focus: {cuda_text}")
def _stop_streaming(self):
"""Stop ZMQ live streaming."""
if self.zmq_client:
self.zmq_client.stop()
self.zmq_client = None
self.stream_out_timestamps.clear()
self.stream_fps_in = 0.0
self.stream_fps_out = 0.0
self.stream_fps_pred = 0.0
self._last_server_frame_id = None
self._last_server_frame_ts = None
self._pending_stream_frame = None
self._pending_stream_detections = None
self._pending_stream_metadata = None
self._last_stream_metadata = {}
self._stream_render_scheduled = False
self.stream_btn.config(text="Connect Live", bg="blue")
self.status_var.set("Streaming stopped")
def toggle_focus_stream(self):
"""Enable/disable focus-score computation for live stream."""
new_state = not self.focus_enabled_var.get()
self.focus_enabled_var.set(new_state)
if self.zmq_client and self.zmq_client.is_streaming():
if new_state:
self.zmq_client.enable_focus()
else:
self.zmq_client.disable_focus()
if new_state:
self.focus_btn.config(text="🎯 Focus ON", bg="#17a2b8")
self.status_var.set("Focus scoring enabled")
else:
self.focus_btn.config(text="🎯 Focus OFF", bg="#6c757d")
self.status_var.set("Focus scoring disabled")
def _on_zmq_frame(self, frame, detections, metadata=None):
"""Callback for new ZMQ frame.
Keep only newest frame and render in UI at a capped rate to stay live."""
now = time.time()
if isinstance(metadata, dict):
source_fps = metadata.get('source_fps')
if isinstance(source_fps, (int, float)):
self.stream_fps_in = float(source_fps)
else:
frame_id = metadata.get('frame_id')
if isinstance(frame_id, int):
if self._last_server_frame_id is not None and self._last_server_frame_ts is not None:
dt = now - self._last_server_frame_ts
df = frame_id - self._last_server_frame_id
if dt > 0 and df >= 0:
self.stream_fps_in = float(df) / float(dt)
self._last_server_frame_id = frame_id
self._last_server_frame_ts = now
prediction_fps = metadata.get('prediction_fps')
if isinstance(prediction_fps, (int, float)):
self.stream_fps_pred = float(prediction_fps)
# Overwrite pending payload: this intentionally drops stale frames.
self._pending_stream_frame = frame
self._pending_stream_detections = detections
self._pending_stream_metadata = metadata
if not self._stream_render_scheduled:
self._stream_render_scheduled = True
self.root.after(0, self._flush_latest_stream_frame)
def _flush_latest_stream_frame(self):
"""Render latest pending stream frame, dropping any intermediate backlog."""
self._stream_render_scheduled = False
if self._pending_stream_frame is None:
return
now = time.time()
if self.stream_ui_max_fps > 0:
min_interval = 1.0 / float(self.stream_ui_max_fps)
dt = now - self._last_ui_render_ts
if dt < min_interval:
delay_ms = max(1, int((min_interval - dt) * 1000))
self._stream_render_scheduled = True
self.root.after(delay_ms, self._flush_latest_stream_frame)
return
frame = self._pending_stream_frame
detections = self._pending_stream_detections if self._pending_stream_detections is not None else []
metadata = self._pending_stream_metadata if isinstance(self._pending_stream_metadata, dict) else {}
self._pending_stream_frame = None
self._pending_stream_detections = None
self._pending_stream_metadata = None
self.current_raw_frame = frame
self.last_results = detections
self._last_stream_metadata = metadata
self.update_canvas(frame, metadata=metadata)
self._last_ui_render_ts = now
# Render FPS is measured on UI flushes, not on received packets.
self.stream_out_timestamps.append(now)
while self.stream_out_timestamps and (now - self.stream_out_timestamps[0]) > 1.0:
self.stream_out_timestamps.popleft()
self.stream_fps_out = float(len(self.stream_out_timestamps))
def _on_focus_update(self, score):
"""Callback for focus score update."""
self.focus_history.append(score)
self.root.after_idle(self.redraw_focus_plot)
# --- Display & Canvas Interaction ---
def regenerate_selected_edge(self):
"""Force regeneration of polygon for selected detection using brush strokes or SAM."""
if self.canvas.selected_obj_idx == -1:
self.status_var.set("No detection selected")
return
if self.canvas.selected_obj_idx >= len(self.last_results):
return
det = self.last_results[self.canvas.selected_obj_idx]
box = det.get('box', det)
# --- NEW: SAM Refinement Path ---
if self.use_sam_refine.get():
try:
self.root.config(cursor="watch")
self.status_var.set("SAM: Refining based on box...")
self.root.update_idletasks()
res = self.inference_client.segment_anything_box(
self.current_raw_frame,
box['x1'], box['y1'], box['x2'], box['y2']
)
poly_pts = res.get('polygon', [])
if poly_pts:
poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
self._set_det_single_polygon(det, poly_arr - [box['x1'], box['y1']])
det.pop('suppress_auto_polygon', None)
self.status_var.set("SAM: Edge refined successfully")
else:
self.status_var.set("SAM: Could not find object in box")
except Exception as e:
self.status_var.set(f"SAM Error: {str(e)}")
finally:
self.root.config(cursor="")
self.redraw_current_frame()
return
# Check if we have manual mask
has_manual_mask = 'manual_mask' in det
print(f"Regenerating edge: has_manual_mask={has_manual_mask}")
if has_manual_mask:
manual_mask = det['manual_mask']
has_bg = np.any(manual_mask == 0)
has_fg = np.any(manual_mask == 1)
print(f" Manual mask: has_bg={has_bg}, has_fg={has_fg}")
# Delete polygon to force regeneration
if 'poly' in det or 'polygons' in det:
det.pop('poly', None)
det.pop('polygons', None)
print(" Deleted existing polygon(s)")
det.pop('suppress_auto_polygon', None)
# Trigger redraw which will regenerate the polygon
self.redraw_current_frame()
# Check if manual mask still exists after redraw
has_manual_mask_after = 'manual_mask' in det
print(f" After redraw: has_manual_mask={has_manual_mask_after}")
if has_manual_mask_after:
self.status_var.set("Edge regenerated - paint more or click again to refine")
else:
self.status_var.set("Edge regenerated")
def clear_brush_strokes(self):
"""Clear brush strokes for selected detection."""
if self.canvas.selected_obj_idx == -1:
self.status_var.set("No detection selected")
return
if self.canvas.selected_obj_idx >= len(self.last_results):
return
self._save_history()
det = self.last_results[self.canvas.selected_obj_idx]
# Clear manual mask and polygon
if 'manual_mask' in det:
del det['manual_mask']
det.pop('poly', None)
det.pop('polygons', None)
self.redraw_current_frame()
self.status_var.set("Brush strokes cleared for selected detection")
def redraw_focus_plot(self):
"""Redraw focus history plot."""
window = self.tk_vars['plot_window_size'].get()
self.ctrl_panel.update_focus_plot(list(self.focus_history), window)
# --- Training & Annotation Submission ---
def _convert_to_yolo_segmentation(self):
"""Convert detections to YOLO segmentation format.
YOLO segmentation format per line:
class_id x1 y1 x2 y2 x3 y3 ... (normalized coordinates 0-1)
"""
if self.current_raw_frame is None or not self.last_results:
return []
img_h, img_w = self.current_raw_frame.shape[:2]
yolo_lines = []
for det in self.last_results:
# Get class
class_name = det.get('name', det.get('label', self._default_class_name()))
class_id = self.class_to_id.get(class_name, 0)
box = det.get('box', det)
polygons = self._det_get_absolute_polygons(det)
if polygons:
for points in polygons:
normalized_points = []
for pt in points:
nx = pt[0] / img_w
ny = pt[1] / img_h
normalized_points.extend([nx, ny])
line = f"{class_id} " + " ".join(f"{coord:.6f}" for coord in normalized_points)
yolo_lines.append(line)
else:
x1, y1 = box.get('x1', 0), box.get('y1', 0)
x2, y2 = box.get('x2', 0), box.get('y2', 0)
cx = ((x1 + x2) / 2) / img_w
cy = ((y1 + y2) / 2) / img_h
w = (x2 - x1) / img_w
h = (y2 - y1) / img_h
line = f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}"
yolo_lines.append(line)
return yolo_lines
def _save_local_annotation_files(self, file_stem, yolo_seg_lines, yolo_det_lines):
"""Persist image + YOLO labels under the configured local output folder."""
root = Path(self.local_output_dir)
targets = {
"segmentation": yolo_seg_lines,
"detection": yolo_det_lines,
}
for task_type, lines in targets.items():
base_dir = root / task_type
img_dir = base_dir / "images"
lbl_dir = base_dir / "labels"
img_dir.mkdir(parents=True, exist_ok=True)
lbl_dir.mkdir(parents=True, exist_ok=True)
cv2.imwrite(str(img_dir / f"{file_stem}.jpg"), self.current_raw_frame)
with open(lbl_dir / f"{file_stem}.txt", "w", encoding="utf-8") as f:
f.write("\n".join(lines) + ("\n" if lines else ""))
classes_txt = root / "classes.txt"
class_names = [name for _, name in sorted((idx, name) for name, idx in self.class_to_id.items())]
classes_txt.write_text("\n".join(class_names) + ("\n" if class_names else ""), encoding="utf-8")
(root / "annotation_mode.txt").write_text(f"{self.annotation_mode}\n", encoding="utf-8")
@staticmethod
def _polygon_area(points):
"""Compute polygon area for COCO annotations."""
arr = np.array(points, dtype=np.float32).reshape(-1, 1, 2)
return float(abs(cv2.contourArea(arr)))
@staticmethod
def _rle_encode_binary_mask(mask):
"""Encode uint8 binary mask into COCO uncompressed RLE."""
pixels = np.asarray(mask, dtype=np.uint8).flatten(order="F")
counts = []
prev = 0
run_len = 0
for pix in pixels:
val = 1 if pix else 0
if val == prev:
run_len += 1
else:
counts.append(run_len)
run_len = 1
prev = val
counts.append(run_len)
return {"size": [int(mask.shape[0]), int(mask.shape[1])], "counts": counts}
@staticmethod
def _rle_decode_binary_mask(seg):
"""Decode COCO uncompressed RLE dict into uint8 binary mask."""
size = seg.get("size", [])
counts = seg.get("counts", [])
if not isinstance(size, list) or len(size) != 2 or not isinstance(counts, list):
return None
h, w = int(size[0]), int(size[1])
flat = np.zeros(h * w, dtype=np.uint8)
idx = 0
val = 0
for run in counts:
run = int(run)
if run <= 0:
val = 1 - val
continue
end = min(idx + run, flat.size)
if val == 1:
flat[idx:end] = 1
idx = end
if idx >= flat.size:
break
val = 1 - val
return flat.reshape((h, w), order="F")
def _build_mask_from_rings(self, img_h, img_w, polygons_abs, holes_abs):
"""Rasterize outer/hole rings to binary mask (outer minus holes)."""
mask = np.zeros((img_h, img_w), dtype=np.uint8)
for poly in polygons_abs:
pts = np.round(poly).astype(np.int32).reshape(-1, 1, 2)
cv2.fillPoly(mask, [pts], 1)
for hole in holes_abs:
pts = np.round(hole).astype(np.int32).reshape(-1, 1, 2)
cv2.fillPoly(mask, [pts], 0)
return mask
def _save_local_coco_annotation_files(self, file_stem):
"""Persist current image+annotations in COCO instance format."""
root = Path(self.local_output_dir) / "coco"
images_dir = root / "images"
images_dir.mkdir(parents=True, exist_ok=True)
image_filename = f"{file_stem}.jpg"
cv2.imwrite(str(images_dir / image_filename), self.current_raw_frame)
coco_path = root / "annotations.json"
if coco_path.exists():
try:
coco = json.loads(coco_path.read_text(encoding="utf-8"))
except Exception:
coco = {}
else:
coco = {}
coco.setdefault("info", {"description": "AareLC local COCO export"})
coco.setdefault("licenses", [])
coco.setdefault("images", [])
coco.setdefault("annotations", [])
coco["categories"] = [
{"id": idx + 1, "name": name}
for name, idx in sorted(self.class_to_id.items(), key=lambda kv: kv[1])
]
coco["meta"] = {"annotation_mode": self.annotation_mode}
img_h, img_w = self.current_raw_frame.shape[:2]
existing_img = next((img for img in coco["images"] if img.get("file_name") == image_filename), None)
if existing_img:
image_id = existing_img["id"]
existing_img.update({"width": int(img_w), "height": int(img_h)})
else:
image_id = (max((img.get("id", 0) for img in coco["images"]), default=0) + 1)
coco["images"].append({
"id": image_id,
"file_name": image_filename,
"width": int(img_w),
"height": int(img_h),
})
coco["annotations"] = [ann for ann in coco["annotations"] if ann.get("image_id") != image_id]
next_ann_id = max((ann.get("id", 0) for ann in coco["annotations"]), default=0) + 1
for det in self.last_results:
class_name = det.get('name', det.get('label', self._default_class_name()))
category_id = self.class_to_id.get(class_name, 0) + 1
box = det.get('box', det)
x1, y1 = float(box['x1']), float(box['y1'])
x2, y2 = float(box['x2']), float(box['y2'])
bw, bh = max(0.0, x2 - x1), max(0.0, y2 - y1)
polygons_abs = self._det_get_absolute_polygons(det)
holes_abs = self._det_get_absolute_holes(det)
if not polygons_abs:
polygons_abs = [np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.float32)]
if holes_abs:
mask = self._build_mask_from_rings(img_h, img_w, polygons_abs, holes_abs)
ys, xs = np.where(mask > 0)
if len(xs) == 0 or len(ys) == 0:
continue
min_x, max_x = float(np.min(xs)), float(np.max(xs))
min_y, max_y = float(np.min(ys)), float(np.max(ys))
segmentation_payload = self._rle_encode_binary_mask(mask)
area_total = float(np.sum(mask > 0))
bbox_payload = [min_x, min_y, max_x - min_x + 1.0, max_y - min_y + 1.0]
else:
segmentation = []
area_total = 0.0
for poly in polygons_abs:
flat = [float(v) for point in poly for v in point]
if len(flat) >= 6:
segmentation.append(flat)
area_total += self._polygon_area(poly)
if not segmentation:
continue
segmentation_payload = segmentation
bbox_payload = [x1, y1, bw, bh]
coco["annotations"].append({
"id": next_ann_id,
"image_id": image_id,
"category_id": category_id,
"segmentation": segmentation_payload,
"area": float(area_total),
"bbox": bbox_payload,
"iscrowd": 0,
})
next_ann_id += 1
coco_path.write_text(json.dumps(coco, indent=2), encoding="utf-8")
(root / "annotation_mode.txt").write_text(f"{self.annotation_mode}\n", encoding="utf-8")
def send_to_training(self):
"""Save YOLO segmentation annotations to database and locally for training."""
if self.review_mode:
self.save_review_annotation()
return
if self.current_raw_frame is None:
messagebox.showwarning("Warning", "No image loaded.")
return
if not self.last_results:
messagebox.showwarning("Warning", "No annotations to save.")
return
try:
img_h, img_w = self.current_raw_frame.shape[:2]
if self._is_coco_multi_mode():
if self.current_local_image_path is not None:
file_stem = self.current_local_image_path.stem
elif self.current_image_id is not None:
file_stem = f"img_{self.current_image_id}"
else:
file_stem = datetime.now().strftime("manual_%Y%m%d_%H%M%S")
confirm_msg = (
f"Save {len(self.last_results)} annotations in COCO multi-polygon mode "
f"to local folder only?\n\nMode: {self._get_annotation_mode_label()}"
)
if not messagebox.askyesno("Confirm", confirm_msg):
return
self.root.config(cursor="watch")
self._save_local_coco_annotation_files(file_stem)
self.status_var.set(f"Saved COCO annotation: {file_stem}")
messagebox.showinfo(
"Success",
(
f"Saved '{file_stem}' in COCO format under:\n{self.local_output_dir / 'coco'}\n\n"
"File: annotations.json (multi-polygon capable)"
),
)
return
if any(self._det_get_holes_relative(det) for det in self.last_results):
if not messagebox.askyesno(
"YOLO Limitation",
(
"Some instances contain hole polygons, but YOLO txt export cannot encode holes.\n"
"Holes will be ignored in YOLO export.\n\nContinue?"
),
):
return
# 1. Prepare YOLO data (Normalized)
segments = []
yolo_seg_lines = [] # For polygons
yolo_det_lines = [] # For bounding boxes
for det in self.last_results:
class_name = det.get('name', det.get('label', self._default_class_name()))
class_id = self.class_to_id.get(class_name, 0)
box = det.get('box', det)
# --- 1a. Generate Detection Line (cx, cy, w, h) ---
x1, y1, x2, y2 = float(box['x1']), float(box['y1']), float(box['x2']), float(box['y2'])
cx = ((x1 + x2) / 2) / img_w
cy = ((y1 + y2) / 2) / img_h
bw = (x2 - x1) / img_w
bh = (y2 - y1) / img_h
yolo_det_lines.append(f"{class_id} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
# --- 1b. Generate Segmentation Line (x1, y1, x2, y2...) ---
polygons = self._det_get_absolute_polygons(det)
if polygons:
for points in polygons:
normalized_points = []
for pt in points:
nx, ny = float(pt[0] / img_w), float(pt[1] / img_h)
normalized_points.extend([nx, ny])
segments.append({
"class_id": class_id,
"points": normalized_points,
"box": [cx, cy, bw, bh]
})
pts_str = " ".join(f"{p:.6f}" for p in normalized_points)
yolo_seg_lines.append(f"{class_id} {pts_str}")
else:
# Fallback: Use box corners as a 4-point polygon
norm_poly = [x1 / img_w, y1 / img_h, x2 / img_w, y1 / img_h, x2 / img_w, y2 / img_h, x1 / img_w,
y2 / img_h]
segments.append({
"class_id": class_id,
"points": norm_poly,
"box": [cx, cy, bw, bh]
})
pts_str = " ".join(f"{p:.6f}" for p in norm_poly)
yolo_seg_lines.append(f"{class_id} {pts_str}")
save_mode = "database + local training sets" if self.current_image_id is not None else "local folder only"
if not messagebox.askyesno("Confirm", f"Save {len(segments)} annotations to {save_mode}?"):
return
self.root.config(cursor="watch")
if self.current_image_id is None:
if self.current_local_image_path is not None:
file_stem = self.current_local_image_path.stem
else:
file_stem = datetime.now().strftime("manual_%Y%m%d_%H%M%S")
self._save_local_annotation_files(file_stem, yolo_seg_lines, yolo_det_lines)
self.status_var.set(f"Saved local annotation: {file_stem}")
messagebox.showinfo(
"Success",
(
f"Saved '{file_stem}' under:\n{self.local_output_dir}\n\n"
"Subfolders created: segmentation/, detection/, classes.txt"
),
)
return
# --- 3. LOCAL SAVING (DB workflow, legacy path layout) ---
file_stem = f"img_{self.current_image_id}"
for task_type in ["seg_training", "det_training"]:
base_dir = Path(f"data/img_sets/{task_type}")
img_dir, lbl_dir = base_dir / "images", base_dir / "labels"
img_dir.mkdir(parents=True, exist_ok=True)
lbl_dir.mkdir(parents=True, exist_ok=True)
# Save the image in both places (required for YOLO training structure)
cv2.imwrite(str(img_dir / f"{file_stem}.jpg"), self.current_raw_frame)
# Save appropriate label format
lines = yolo_seg_lines if "seg" in task_type else yolo_det_lines
with open(lbl_dir / f"{file_stem}.txt", "w") as f:
f.write("\n".join(lines))
# --- 4. DATABASE SAVING ---
db_result = self.db_client.save_annotation(
self.current_image_id,
segments,
username=self.pref_username.get()
)
status_msg = f"Saved locally ({file_stem}) and to DB (ID: {db_result.get('annotation_id')})"
self.status_var.set((status_msg[:120] + '...') if len(status_msg) > 123 else status_msg)
# FIX: Removed the undefined 'img_path' variable and updated the message
messagebox.showinfo(
"Success",
f"Annotated image '{file_stem}' saved to local training sets (Detection & Segmentation) and synced to database."
)
# 5. Fetch next image automatically
self.fetch_db_image()
except Exception as e:
messagebox.showerror("Error", f"Failed to save:\n{str(e)}")
import traceback; traceback.print_exc()
finally:
self.root.config(cursor="")
def skip_current_image(self):
"""Skip the current image and fetch the next one."""
if self.current_image_id is None:
messagebox.showwarning(
"Warning",
"No image to skip. Please use 'Next Image' button first."
)
return
try:
# Confirm skip
response = messagebox.askyesno(
"Confirm Skip",
f"Skip image ID {self.current_image_id}?\n\n"
"This image will be marked as skipped in the database."
)
if not response:
return
self.root.config(cursor="watch")
result = self.db_client.skip_image(self.current_image_id)
self.status_var.set(f"Skipped image ID {self.current_image_id}")
# Fetch next image automatically
self.fetch_db_image()
except Exception as e:
messagebox.showerror("Error", f"Failed to skip image:\n{str(e)}")
self.status_var.set("Skip failed")
import traceback
traceback.print_exc()
finally:
self.root.config(cursor="")
def show_sync_dialog(self):
"""Show dialog to sync missing local images and labels from DB."""
dialog = tk.Toplevel(self.root)
dialog.title("Sync DB -> Local")
dialog.geometry("640x420")
dialog.minsize(520, 340)
dialog.resizable(True, True)
dialog.transient(self.root)
dialog.grab_set()
frame = tk.Frame(dialog, padx=20, pady=20)
frame.pack(fill=tk.BOTH, expand=True)
tk.Label(
frame,
text="Copy only missing local image/annotation files from DB records.",
font=('Arial', 10),
).grid(row=0, column=0, columnspan=2, sticky=tk.W, pady=(0, 12))
tk.Label(frame, text="Since (YYYY-MM-DD or ISO):").grid(
row=1, column=0, sticky=tk.W, pady=6
)
tk.Entry(frame, textvariable=self.sync_since, width=32).grid(
row=1, column=1, sticky=tk.W, padx=8
)
tk.Label(frame, text="Until (YYYY-MM-DD or ISO):").grid(
row=2, column=0, sticky=tk.W, pady=6
)
tk.Entry(frame, textvariable=self.sync_until, width=32).grid(
row=2, column=1, sticky=tk.W, padx=8
)
tk.Label(frame, text="Limit:").grid(row=3, column=0, sticky=tk.W, pady=6)
tk.Entry(frame, textvariable=self.sync_limit, width=12).grid(
row=3, column=1, sticky=tk.W, padx=8
)
tk.Checkbutton(
frame,
text="Include unannotated images",
variable=self.sync_include_unannotated,
).grid(row=4, column=0, columnspan=2, sticky=tk.W, pady=(10, 4))
tk.Checkbutton(
frame,
text="Dry run (preview only, no files written)",
variable=self.sync_dry_run,
).grid(row=5, column=0, columnspan=2, sticky=tk.W, pady=4)
hint = (
"Tip: run dry-run first, then uncheck dry-run to copy files."
)
tk.Label(frame, text=hint, fg="#555", font=('Arial', 9)).grid(
row=6, column=0, columnspan=2, sticky=tk.W, pady=(10, 4)
)
def run_sync():
try:
limit = int(self.sync_limit.get())
if limit <= 0:
raise ValueError("Limit must be > 0")
self.root.config(cursor="watch")
dialog.config(cursor="watch")
self.status_var.set("Syncing local files from database...")
self.root.update_idletasks()
result = self.db_client.sync_local_annotations(
since=self.sync_since.get().strip() or None,
until=self.sync_until.get().strip() or None,
include_unannotated=self.sync_include_unannotated.get(),
dry_run=self.sync_dry_run.get(),
limit=limit,
)
stats = result.get("stats", {})
failures = result.get("failures", [])
summary = (
f"Scanned: {stats.get('scanned', 0)}\n"
f"Images copied: {stats.get('image_copied', 0)}\n"
f"Images existing: {stats.get('image_existing', 0)}\n"
f"Images failed: {stats.get('image_failed', 0)}\n"
f"Labels copied: {stats.get('label_copied', 0)}\n"
f"Labels existing: {stats.get('label_existing', 0)}\n"
f"Labels failed: {stats.get('label_failed', 0)}"
)
if failures:
sample_lines = [f[:80] + "..." if len(f) > 80 else f for f in failures[:10]]
sample = "\n".join(sample_lines)
messagebox.showwarning(
"Sync Completed with Warnings",
f"{summary}\n\nFirst failures:\n{sample}",
)
else:
messagebox.showinfo("Sync Completed", summary)
self.status_var.set("DB->local sync completed")
except Exception as e:
messagebox.showerror("Sync Error", str(e))
self.status_var.set("DB->local sync failed")
finally:
self.root.config(cursor="")
dialog.config(cursor="")
button_row = tk.Frame(frame)
button_row.grid(row=7, column=0, columnspan=2, sticky=tk.EW, pady=(18, 0))
tk.Button(button_row, text="Close", command=dialog.destroy, width=12).pack(
side=tk.LEFT, padx=8
)
tk.Button(
button_row,
text="Run Sync",
command=run_sync,
bg="#007bff",
fg="white",
width=12,
font=('Arial', 10, 'bold'),
).pack(side=tk.LEFT, padx=8)
def show_zmq_inference_control_dialog(self):
"""Control runtime settings of remote zmq_inference server via HTTP API."""
dialog = tk.Toplevel(self.root)
dialog.title("ZMQ Inference Control")
dialog.geometry("900x780")
dialog.minsize(760, 700)
dialog.resizable(True, True)
dialog.transient(self.root)
dialog.grab_set()
frame = tk.Frame(dialog, padx=16, pady=14)
frame.pack(fill=tk.BOTH, expand=True)
frame.grid_columnconfigure(1, weight=1)
control_url_var = tk.StringVar(value=self.pref_zmq_control_url.get().strip() or "http://localhost:8090")
control_token_var = tk.StringVar(value=self.pref_zmq_control_token.get().strip())
conf_var = tk.DoubleVar(value=0.25)
infer_scale_var = tk.DoubleVar(value=1.0)
skip_var = tk.IntVar(value=0)
max_fps_var = tk.DoubleVar(value=0.0)
device_var = tk.StringVar(value="0")
shard_count_var = tk.IntVar(value=1)
shard_index_var = tk.IntVar(value=0)
model_choice_var = tk.StringVar(value="")
model_seg_hint_var = tk.StringVar(value="Segmentation capable: unknown")
cuda_hint_var = tk.StringVar(value="Server CUDA: unknown")
task_var = tk.StringVar(value="auto")
tracker_var = tk.StringVar(value="none")
bayer_flip_var = tk.BooleanVar(value=False)
publish_enabled_var = tk.BooleanVar(value=True)
compute_target_point_var = tk.BooleanVar(value=False)
focus_enabled_var = tk.BooleanVar(value=False)
focus_epics_enabled_var = tk.BooleanVar(value=False)
focus_pv_var = tk.StringVar(value="X10SA-ES-MS:cam1:FocusScore")
focus_every_var = tk.IntVar(value=3)
focus_scale_var = tk.DoubleVar(value=0.5)
focus_pv_min_period_var = tk.DoubleVar(value=100.0)
model_items_by_label = {}
server_cuda_available = {"value": False}
model_refresh_job = {"id": None}
runtime_model_paths = {"engine": "", "pt": ""}
tk.Label(frame, text="Control URL:", font=('Arial', 10, 'bold')).grid(row=0, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=control_url_var, width=52).grid(row=0, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Control Token:").grid(row=1, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=control_token_var, width=30, show="*").grid(row=1, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Confidence:").grid(row=2, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=conf_var, width=16).grid(row=2, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Infer Scale:").grid(row=3, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=infer_scale_var, width=16).grid(row=3, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Skip Frames:").grid(row=4, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=skip_var, width=16).grid(row=4, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Max Predict FPS (0=off):").grid(row=5, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=max_fps_var, width=16).grid(row=5, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Model (models/):").grid(row=6, column=0, sticky=tk.W, pady=4)
model_combo = ttk.Combobox(frame, textvariable=model_choice_var, width=58, state="readonly")
model_combo.grid(row=6, column=1, sticky=tk.W, padx=8)
tk.Label(frame, textvariable=model_seg_hint_var, fg="#126b2c", font=('Arial', 9, 'bold')).grid(
row=7, column=1, sticky=tk.W, padx=8, pady=(0, 2)
)
tk.Label(frame, textvariable=cuda_hint_var, fg="#1c4f8a", font=('Arial', 9, 'bold')).grid(
row=8, column=1, sticky=tk.W, padx=8, pady=(0, 2)
)
tk.Label(frame, text="Task:").grid(row=9, column=0, sticky=tk.W, pady=4)
task_combo = ttk.Combobox(
frame,
textvariable=task_var,
values=("auto", "detect", "segment"),
width=14,
state="readonly",
)
task_combo.grid(row=9, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Tracker:").grid(row=10, column=0, sticky=tk.W, pady=4)
tracker_combo = ttk.Combobox(
frame,
textvariable=tracker_var,
values=("none", "bytetrack", "botsort"),
width=14,
state="readonly",
)
tracker_combo.grid(row=10, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Device (cpu or GPU idx):").grid(row=11, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=device_var, width=16).grid(row=11, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Shard Count (1=off, 0=auto):").grid(row=12, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=shard_count_var, width=16).grid(row=12, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Shard Index:").grid(row=13, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=shard_index_var, width=16).grid(row=13, column=1, sticky=tk.W, padx=8)
tk.Checkbutton(frame, text="Bayer Flip", variable=bayer_flip_var).grid(
row=14, column=1, sticky=tk.W, padx=8, pady=(6, 2)
)
tk.Checkbutton(frame, text="Publish Output", variable=publish_enabled_var).grid(
row=15, column=1, sticky=tk.W, padx=8, pady=(2, 2)
)
tk.Checkbutton(frame, text="Compute Target Point", variable=compute_target_point_var).grid(
row=16, column=1, sticky=tk.W, padx=8, pady=(2, 2)
)
tk.Checkbutton(frame, text="Compute Focus Score", variable=focus_enabled_var).grid(
row=17, column=1, sticky=tk.W, padx=8, pady=(2, 2)
)
tk.Checkbutton(frame, text="Write Focus to EPICS", variable=focus_epics_enabled_var).grid(
row=18, column=1, sticky=tk.W, padx=8, pady=(2, 2)
)
tk.Label(frame, text="Focus EPICS PV:").grid(row=19, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=focus_pv_var, width=52).grid(row=19, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Focus Every N Frames:").grid(row=20, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=focus_every_var, width=16).grid(row=20, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="Focus Scale (0-1]:").grid(row=21, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=focus_scale_var, width=16).grid(row=21, column=1, sticky=tk.W, padx=8)
tk.Label(frame, text="EPICS Min Period (ms):").grid(row=22, column=0, sticky=tk.W, pady=4)
tk.Entry(frame, textvariable=focus_pv_min_period_var, width=16).grid(row=22, column=1, sticky=tk.W, padx=8)
hint = (
"Apply changes to the running zmq_inference service.\n"
"Model/task changes trigger hot reload for all downstream ZMQ consumers."
)
tk.Label(frame, text=hint, fg="#555", justify=tk.LEFT, font=('Arial', 9)).grid(
row=23, column=0, columnspan=2, sticky=tk.W, pady=(10, 6)
)
def _base_url():
base = control_url_var.get().strip().rstrip("/")
if not base:
raise ValueError("Control URL is empty")
return base
def _norm_path(value):
if not value:
return ""
return os.path.normpath(os.path.expanduser(str(value).strip()))
def _canonical_model_ref(value):
"""Normalize model path refs to improve matching across server/client cwd differences."""
raw = str(value or "").strip()
if not raw:
return ""
normalized = raw.replace("\\", "/")
if "/models/" in normalized:
return "models/" + normalized.split("/models/", 1)[1]
if normalized.startswith("../models/"):
return "models/" + normalized[len("../models/"):]
if normalized.startswith("./models/"):
return "models/" + normalized[len("./models/"):]
return normalized
def _set_segmentation_hint(item):
if not item:
model_seg_hint_var.set("Segmentation capable: unknown")
return
yes_no = "yes" if bool(item.get("segmentation_capable")) else "no"
model_seg_hint_var.set(f"Segmentation capable: {yes_no}")
def _model_label(item):
name = str(item.get("name") or item.get("project_path") or item.get("path") or "model")
kind = str(item.get("kind", "")).upper()
seg = "SEG" if bool(item.get("segmentation_capable")) else "DET"
kind_lower = str(item.get("kind", "")).lower()
cuda_mark = ""
if server_cuda_available["value"] and kind_lower in ("engine", "engines"):
cuda_mark = " | CUDA"
if kind:
return f"{name} [{kind} | {seg}{cuda_mark}]"
return f"{name} [{seg}]"
def _match_model_item(model_path):
if not model_path:
return None
normalized = _norm_path(_canonical_model_ref(model_path))
for item in model_items_by_label.values():
candidates = {
_norm_path(_canonical_model_ref(item.get("path"))),
_norm_path(_canonical_model_ref(item.get("project_path"))),
_norm_path(_canonical_model_ref(item.get("name"))),
}
if normalized in candidates:
return item
return None
def _apply_model_item(item, update_task=True):
if not item:
return
if update_task:
task_hint = str(item.get("task_hint") or "").lower()
if task_hint in ("detect", "segment"):
task_var.set(task_hint)
_set_segmentation_hint(item)
def refresh_models(quiet=False):
try:
base = _base_url()
token = control_token_var.get().strip()
headers = {"X-API-Key": token} if token else None
resp = requests.get(f"{base}/models", headers=headers, timeout=3.0)
resp.raise_for_status()
payload = resp.json()
models = payload.get("models", [])
caps = payload.get("capabilities", {})
server_cuda_available["value"] = bool(caps.get("cuda_available", False))
cuda_hint_var.set(f"Server CUDA: {'yes' if server_cuda_available['value'] else 'no'}")
except Exception:
# Fallback for older servers without /models endpoint.
local_models = []
models_root = Path(__file__).resolve().parent / "models"
if models_root.exists():
for p in sorted(models_root.rglob("*")):
if not p.is_file() or p.suffix.lower() not in {".pt", ".engine", ".engines", ".yaml", ".yml", ".onnx"}:
continue
rel = str(p.relative_to(models_root))
seg = any(tok in p.name.lower() for tok in ("seg", "segment", "mask"))
local_models.append(
{
"name": rel,
"path": str(p),
"project_path": str(Path("models") / rel),
"kind": p.suffix.lower().lstrip("."),
"segmentation_capable": seg,
"task_hint": "segment" if seg else "detect",
}
)
models = local_models
if not quiet and not models:
messagebox.showwarning(
"Models Unavailable",
"Could not fetch models from server and no local models folder entries were found.",
)
# Ensure current runtime model is selectable even if it is outside scanned models/.
current_model_path = runtime_model_paths.get("engine", "").strip() or runtime_model_paths.get("pt", "").strip()
if current_model_path:
current_norm = _norm_path(_canonical_model_ref(current_model_path))
has_match = False
for item in models:
candidates = {
_norm_path(_canonical_model_ref(item.get("path"))),
_norm_path(_canonical_model_ref(item.get("project_path"))),
_norm_path(_canonical_model_ref(item.get("name"))),
}
if current_norm in candidates:
has_match = True
break
if not has_match:
normalized_current = _canonical_model_ref(current_model_path)
current_name = normalized_current.split("/")[-1] if "/" in normalized_current else normalized_current
current_kind = "engine" if normalized_current.endswith((".engine", ".engines")) else "pt"
seg_cap = (
str(task_var.get()).strip().lower() == "segment"
or any(tok in current_name.lower() for tok in ("seg", "segment", "mask"))
)
models.insert(
0,
{
"name": current_name,
"path": current_model_path,
"project_path": normalized_current,
"kind": current_kind,
"segmentation_capable": seg_cap,
"task_hint": "segment" if seg_cap else "detect",
},
)
labels = []
model_items_by_label.clear()
for item in models:
label = _model_label(item)
if len(label) > 60:
label = label[:57] + "..."
labels.append(label)
model_items_by_label[label] = item
model_combo["values"] = labels
if not labels:
model_choice_var.set("")
_set_segmentation_hint(None)
return
current_model_path = runtime_model_paths.get("engine", "").strip() or runtime_model_paths.get("pt", "").strip()
selected_item = _match_model_item(current_model_path)
if selected_item:
selected_label = _model_label(selected_item)
model_choice_var.set(selected_label)
_set_segmentation_hint(selected_item)
else:
model_choice_var.set(labels[0])
_set_segmentation_hint(model_items_by_label.get(labels[0]))
def _on_model_selected(_event=None):
item = model_items_by_label.get(model_choice_var.get())
_apply_model_item(item, update_task=True)
model_combo.bind("<<ComboboxSelected>>", _on_model_selected)
def load_current(quiet=False):
try:
base = _base_url()
token = control_token_var.get().strip()
headers = {"X-API-Key": token} if token else None
self.root.config(cursor="watch")
dialog.config(cursor="watch")
resp = requests.get(f"{base}/config", headers=headers, timeout=3.0)
resp.raise_for_status()
data = resp.json()
cfg = data.get("config", data)
caps = data.get("capabilities", {})
if "cuda_available" in caps:
server_cuda_available["value"] = bool(caps.get("cuda_available"))
cuda_hint_var.set(f"Server CUDA: {'yes' if server_cuda_available['value'] else 'no'}")
conf_var.set(float(cfg.get("conf", conf_var.get())))
infer_scale_var.set(float(cfg.get("infer_scale", infer_scale_var.get())))
skip_var.set(int(cfg.get("skip", skip_var.get())))
max_fps_var.set(float(cfg.get("max_fps", max_fps_var.get())))
device_var.set(str(cfg.get("device", device_var.get())))
shard_count_var.set(int(cfg.get("shard_count", shard_count_var.get())))
shard_index_var.set(int(cfg.get("shard_index", shard_index_var.get())))
task_var.set(str(cfg.get("task", task_var.get())))
tracker_var.set(str(cfg.get("tracker", tracker_var.get())))
runtime_model_paths["engine"] = str(cfg.get("engine", "")).strip()
runtime_model_paths["pt"] = str(cfg.get("pt", "")).strip()
bayer_flip_var.set(bool(cfg.get("bayer_flip", bayer_flip_var.get())))
publish_enabled_var.set(bool(cfg.get("publish_enabled", publish_enabled_var.get())))
compute_target_point_var.set(bool(cfg.get("compute_target_point", compute_target_point_var.get())))
focus_enabled_var.set(bool(cfg.get("focus_enabled", focus_enabled_var.get())))
focus_epics_enabled_var.set(bool(cfg.get("focus_epics_enabled", focus_epics_enabled_var.get())))
focus_pv_var.set(str(cfg.get("focus_pv", focus_pv_var.get())))
focus_every_var.set(int(cfg.get("focus_every", focus_every_var.get())))
focus_scale_var.set(float(cfg.get("focus_scale", focus_scale_var.get())))
focus_pv_min_period_var.set(float(cfg.get("focus_pv_min_period_ms", focus_pv_min_period_var.get())))
refresh_models(quiet=True)
current_item = _match_model_item(runtime_model_paths.get("engine") or runtime_model_paths.get("pt"))
if current_item:
model_choice_var.set(_model_label(current_item))
_set_segmentation_hint(current_item)
else:
_set_segmentation_hint(None)
self.pref_zmq_control_url.set(base)
self.pref_zmq_control_token.set(token)
self.status_var.set("Loaded ZMQ inference runtime config")
except Exception as e:
if not quiet:
messagebox.showerror("Load Error", f"Failed loading config:\n{e}")
self.status_var.set("Failed loading ZMQ inference config")
finally:
self.root.config(cursor="")
dialog.config(cursor="")
def _schedule_refresh_on_url_change(*_args):
if model_refresh_job["id"] is not None:
try:
dialog.after_cancel(model_refresh_job["id"])
except Exception:
pass
def _run():
model_refresh_job["id"] = None
try:
load_current(quiet=True)
except Exception:
pass
model_refresh_job["id"] = dialog.after(500, _run)
control_url_var.trace_add("write", _schedule_refresh_on_url_change)
def apply_changes():
try:
base = _base_url()
token = control_token_var.get().strip()
headers = {"X-API-Key": token} if token else None
selected_item = model_items_by_label.get(model_choice_var.get())
if not selected_item:
raise ValueError("Select a model from models/ before applying.")
selected_kind = str(selected_item.get("kind") or "").lower()
selected_project_path = str(
selected_item.get("project_path") or selected_item.get("name") or ""
).strip()
if not selected_project_path:
raise ValueError("Selected model has no usable models/ path.")
selected_engine = selected_project_path if selected_kind in ("engine", "engines") else ""
selected_pt = selected_project_path if selected_kind not in ("engine", "engines") else ""
device_value = str(device_var.get()).strip()
if not device_value:
raise ValueError("Device is required (cpu or CUDA index).")
shard_count_value = int(shard_count_var.get())
shard_index_value = int(shard_index_var.get())
payload = {
"conf": float(conf_var.get()),
"infer_scale": float(infer_scale_var.get()),
"skip": int(skip_var.get()),
"max_fps": float(max_fps_var.get()),
"device": device_value,
"shard_count": shard_count_value,
"shard_index": shard_index_value,
"task": str(task_var.get()),
"tracker": str(tracker_var.get()),
"bayer_flip": bool(bayer_flip_var.get()),
"publish_enabled": bool(publish_enabled_var.get()),
"compute_target_point": bool(compute_target_point_var.get()),
"focus_enabled": bool(focus_enabled_var.get()),
"focus_epics_enabled": bool(focus_epics_enabled_var.get()),
"focus_pv": str(focus_pv_var.get()).strip(),
"focus_every": int(focus_every_var.get()),
"focus_scale": float(focus_scale_var.get()),
"focus_pv_min_period_ms": float(focus_pv_min_period_var.get()),
"engine": selected_engine,
"pt": selected_pt,
}
self.root.config(cursor="watch")
dialog.config(cursor="watch")
resp = requests.post(f"{base}/config", json=payload, headers=headers, timeout=6.0)
if resp.status_code >= 400:
raise RuntimeError(f"HTTP {resp.status_code}: {resp.text}")
out = resp.json()
self.pref_zmq_control_url.set(base)
self.pref_zmq_control_token.set(token)
messagebox.showinfo(
"Config Applied",
f"Updated runtime config.\nModel reloaded: {bool(out.get('model_reloaded', False))}",
)
self.status_var.set("Applied ZMQ inference runtime config")
except Exception as e:
messagebox.showerror("Apply Error", f"Failed applying config:\n{e}")
self.status_var.set("Failed applying ZMQ inference config")
finally:
self.root.config(cursor="")
dialog.config(cursor="")
button_row = tk.Frame(frame)
button_row.grid(row=24, column=0, columnspan=2, sticky=tk.EW, pady=(12, 0))
tk.Button(button_row, text="Close", command=dialog.destroy, width=12).pack(side=tk.LEFT, padx=6)
tk.Button(button_row, text="Refresh Models", command=refresh_models, width=14).pack(side=tk.LEFT, padx=6)
tk.Button(button_row, text="Load Current", command=load_current, width=14).pack(side=tk.LEFT, padx=6)
tk.Button(
button_row,
text="Apply",
command=apply_changes,
bg="#28a745",
fg="white",
width=12,
font=('Arial', 10, 'bold'),
).pack(side=tk.LEFT, padx=6)
# Initialize from the currently selected server.
load_current(quiet=True)
# --- Preferences ---
def show_preferences(self):
"""Show a single consolidated preferences dialog."""
dialog = tk.Toplevel(self.root)
dialog.title("Server Settings")
dialog.geometry("640x500")
dialog.minsize(500, 390)
dialog.resizable(True, True)
dialog.transient(self.root)
dialog.grab_set()
main_frame = tk.Frame(dialog, padx=20, pady=20)
main_frame.pack(fill=tk.BOTH, expand=True)
original_stream_url = self.zmq_det_url.get().strip()
# Define fields to display
config_fields = [
("Inference API URL:", self.pref_inference_url, False),
("Database API URL:", self.pref_db_url, False),
("Image Download URL:", self.pref_download_url, False),
("ZMQ Control URL:", self.pref_zmq_control_url, False),
("ZMQ Control Token:", self.pref_zmq_control_token, True),
("ZMQ Unified Stream:", self.zmq_det_url, False),
("Shared Password:", self.pref_shared_pw, True),
("Annotator Name:", self.pref_username, False)
]
combo_values_by_label = {
"Image Download URL:": self.pref_download_url_options,
"ZMQ Control URL:": self.pref_zmq_control_options,
"ZMQ Unified Stream:": self.pref_zmq_stream_options,
}
for i, (label_text, var, is_pwd) in enumerate(config_fields):
tk.Label(main_frame, text=label_text, font=('Arial', 10)).grid(row=i, column=0, sticky=tk.W, pady=8)
if label_text in combo_values_by_label and not is_pwd:
entry = ttk.Combobox(main_frame, textvariable=var, values=combo_values_by_label[label_text], width=38)
else:
entry = tk.Entry(main_frame, textvariable=var, width=40, show="*" if is_pwd else "")
entry.grid(row=i, column=1, sticky=tk.EW, padx=10)
def apply_and_close():
# Update the actual clients with the new values
self.inference_client.set_server_url(self.pref_inference_url.get())
self.db_client.set_server_url(self.pref_db_url.get())
self.db_client.set_shared_password(self.pref_shared_pw.get())
if hasattr(self.db_client, 'set_download_base'):
self.db_client.set_download_base(self.pref_download_url.get())
new_stream_url = self.zmq_det_url.get().strip()
if new_stream_url and new_stream_url != original_stream_url:
try:
parsed = urllib.parse.urlparse(new_stream_url)
host = parsed.hostname
if host:
inferred_control = f"http://{host}:8090"
self.pref_zmq_control_url.set(inferred_control)
except Exception:
pass
self.status_var.set("Settings updated successfully.")
dialog.destroy()
btn_frame = tk.Frame(main_frame)
btn_frame.grid(row=len(config_fields), column=0, columnspan=2, pady=20)
tk.Button(btn_frame, text="Cancel", command=dialog.destroy, width=12).pack(side=tk.LEFT, padx=10)
tk.Button(
btn_frame,
text="Save Changes",
command=apply_and_close,
bg="#28a745",
fg="white",
width=12,
font=('Arial', 10, 'bold')
).pack(side=tk.LEFT, padx=10)
if __name__ == "__main__":
# --- X11 CRASH PREVENTATIVE MEASURES ---
# Force legacy rendering to avoid RenderAddGlyphs BadLength errors on virtual displays
os.environ["GDK_RENDERING"] = "image"
os.environ["XFT_CONFIG"] = "/dev/null"
os.environ["XLIB_SKIP_ARGB_VISUALS"] = "1"
try:
root = tk.Tk()
# Disable X Input Methods and Anti-aliasing which trigger the BadLength error
root.tk.call('tk', 'useinputmethods', '0')
# Force standard core bitmap fonts for ALL widgets globally
# This bypasses the RENDER extension entirely
root.option_add("*Font", "fixed")
root.option_add("*TLabel*Font", "fixed")
root.option_add("*TButton*Font", "fixed")
root.option_add("*Menu*Font", "fixed")
gui = InferenceGUI(root)
root.mainloop()
except tk.TclError as e:
print(f"CRITICAL: Could not connect to X11 display. Check your SSH -X forwarding.")
print(f"Error: {e}")
Reference in New Issue View Git Blame Copy Permalink