migration and splitting AareLC

src/clip/bl_clip_manager.py

@@ -0,0 +1,701 @@
+import argparse
+import datetime
+import json
+import os
+import time
+from typing import Any
+
+import requests
+import urllib3
+from websocket import create_connection
+
+# Disable SSL warnings if using local IP without a certificate
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+# --- CONFIGURATION ---
+NAS_IP = os.getenv("NAS_IP")  # Synology IP
+NAS_USER = os.getenv("NAS_USER")
+NAS_PASS = os.getenv("NAS_PASS")
+AAREDB_SHARED_PASSWORD = os.getenv("AAREDB_SHARED_PASSWORD")
+
+# Where downloaded files are stored by this script runtime
+DEFAULT_OUTPUT_DIR = os.getenv("CLIP_OUTPUT_DIR", "./clips")
+CLIP_TMP_DIR = os.getenv("CLIP_TMP_DIR", "/sls/mx/misc")
+CLIP_PUBLIC_BASE = os.getenv("CLIP_PUBLIC_BASE", "").strip()
+NAS_EXPORT_DSTDIR = os.getenv("NAS_EXPORT_DSTDIR", "EventClips")
+
+# Dispatcher endpoints for clip jobs
+DISPATCHER_WS_URL = os.getenv(
+    "DISPATCHER_WS_URL",
+    "wss://mx-aaredb-dmz-01.psi.ch/dispatcher/protected_router/clip_runner/ws/jobs",
+    #"wss://127.0.0.1:8001/dispatcher/protected_router/clip_runner/ws/jobs",
+
+)
+DISPATCHER_STATUS_URL = os.getenv(
+    "DISPATCHER_STATUS_URL",
+    "https://mx-aaredb-dmz-01.psi.ch/dispatcher/protected_router/clip_runner/jobs/update_status",
+    #"https://127.0.0.1:8001/dispatcher/protected_router/clip_runner/jobs/update_status",
+
+)
+DISPATCHER_CA_CERT = os.getenv("DISPATCHER_CA_CERT", "").strip() or None
+DISPATCHER_CLIENT_CERT = os.getenv("DISPATCHER_CLIENT_CERT", "").strip() or None
+DISPATCHER_CLIENT_KEY = os.getenv("DISPATCHER_CLIENT_KEY", "").strip() or None
+# If a CA cert path is given, use it for verification; otherwise fall back to bool flag
+DISPATCHER_VERIFY_SSL: bool | str = DISPATCHER_CA_CERT or (
+    os.getenv("DISPATCHER_VERIFY_SSL", "false").lower() in ("1", "true", "yes")
+)
+
+
+def ensure_env_ready():
+    missing = []
+    for key, value in (
+        ("NAS_IP", NAS_IP),
+        ("NAS_USER", NAS_USER),
+        ("NAS_PASS", NAS_PASS),
+        ("AAREDB_SHARED_PASSWORD", AAREDB_SHARED_PASSWORD),
+    ):
+        if not value:
+            missing.append(key)
+    if missing:
+        raise RuntimeError(f"Missing required environment variables: {', '.join(missing)}")
+
+
+def get_sid(session: str = "SurveillanceStation", request_token: bool = False) -> tuple[str, str | None]:
+    """Logs in and returns (sid, synotoken). synotoken is required for POST requests on DSM 7+."""
+    url = f"http://{NAS_IP}:5000/webapi/auth.cgi"
+    params = {
+        "api": "SYNO.API.Auth",
+        "version": "3",
+        "method": "login",
+        "account": NAS_USER,
+        "passwd": NAS_PASS,
+        "session": session,
+        "format": "sid",
+    }
+    if request_token:
+        params["enable_syno_token"] = "yes"
+    r = requests.get(url, params=params, timeout=30)
+    r.raise_for_status()
+    payload = r.json()
+    if not payload.get("success"):
+        raise RuntimeError(f"NAS login failed: {payload}")
+    data = payload["data"]
+    return data["sid"], data.get("synotoken")
+
+
+def logout_sid(sid: str):
+    url = (
+        f"http://{NAS_IP}:5000/webapi/auth.cgi"
+        f"?api=SYNO.API.Auth&version=3&method=logout&sid={sid}"
+    )
+    try:
+        requests.get(url, timeout=15)
+    except Exception:
+        pass
+
+
+def start_range_export(sid: str, camera_id: int, start_time: int, end_time: int, event_name: str):
+    """
+    Starts a time-range export task and returns its download id (dlid).
+    """
+    url = f"http://{NAS_IP}:5000/webapi/entry.cgi"
+    params = {
+        "api": "SYNO.SurveillanceStation.Recording",
+        "method": "RangeExport",
+        "version": "6",
+        "_sid": sid,
+        "camId": camera_id,
+        "fromTime": start_time,
+        "toTime": end_time,
+        "fileName": event_name,
+    }
+    r = requests.get(url, params=params, timeout=60)
+    r.raise_for_status()
+    return r.json()
+
+
+def start_nas_save_export(
+    sid: str,
+    camera_id: int,
+    start_time: int,
+    end_time: int,
+    event_name: str,
+    nas_dstdir: str,
+):
+    """
+    Ask Surveillance Station to save export directly on the NAS (no local download).
+    """
+    url = f"http://{NAS_IP}:5000/webapi/entry.cgi"
+    # Pre-check export feasibility in the camelCase flavor used by SS docs examples.
+    check_params = {
+        "api": "SYNO.SurveillanceStation.Recording.Export",
+        "method": "CheckAvailableExport",
+        "version": "1",
+        "_sid": sid,
+        "camIdList": str(camera_id),
+        "startTime": start_time,
+        "stopTime": end_time,
+        "freeSize": 0,
+    }
+    check_payload = None
+    try:
+        chk = requests.get(url, params=check_params, timeout=60)
+        chk.raise_for_status()
+        check_payload = chk.json()
+        print(f"CheckAvailableExport: {check_payload}")
+    except Exception as exc:
+        print(f"CheckAvailableExport failed (continuing with Save attempts): {exc}")
+
+        # Try multiple variants because SS builds differ significantly in parameter naming
+        save_variants = [
+            {
+                "api": "SYNO.SurveillanceStation.Recording.Export",
+                "method": "Save",
+                "version": "1",
+                "_sid": sid,
+                "name": event_name,
+                "srcDsId": 0,
+                "dstDsId": 0,
+                "dstdir": nas_dstdir,
+                "startTime": start_time,
+                "stopTime": end_time,
+                "isoverwrite": 1,
+                "camIdList": str(camera_id),
+            },
+            {
+                "api": "SYNO.SurveillanceStation.Recording.Export",
+                "method": "Save",
+                "version": "1",
+                "_sid": sid,
+                "name": f"\"{event_name}\"",
+                "srcDsId": 0,
+                "dstDsId": 0,
+                "dstdir": f"\"{nas_dstdir}\"",
+                "start_time": start_time,
+                "stop_time": end_time,
+                "isoverwrite": 1,
+                "camlistid": f"\"{camera_id}\"",
+            },
+        ]
+
+        last_payload = None
+    for idx, params in enumerate(save_variants, start=1):
+        r = requests.get(url, params=params, timeout=60)
+        r.raise_for_status()
+        payload = r.json()
+        print(f"Save variant #{idx}: {payload}")
+        if payload.get("success"):
+            return payload, check_payload
+        last_payload = payload
+
+    return (
+        last_payload or {"success": False, "error": {"code": -1, "message": "Save export failed"}},
+        check_payload,
+    )
+
+
+def wait_for_export(sid: str, dlid: int, timeout_sec: int = 600, poll_interval_sec: int = 2):
+    """
+    Polls export progress and returns the final file extension.
+    """
+    url = f"http://{NAS_IP}:5000/webapi/entry.cgi"
+    deadline = time.time() + timeout_sec
+
+    while time.time() < deadline:
+        params = {
+            "api": "SYNO.SurveillanceStation.Recording",
+            "method": "GetRangeExportProgress",
+            "version": "6",
+            "_sid": sid,
+            "dlid": dlid,
+        }
+        r = requests.get(url, params=params, timeout=60)
+        r.raise_for_status()
+        result = r.json()
+        if not result.get("success"):
+            return None, result
+
+        data = result.get("data", {})
+        progress = data.get("progress", 0)
+        print(f"Export progress (dlid={dlid}): {progress}%")
+
+        if progress == 100:
+            return data.get("fileExt") or "mp4", result
+        if progress == -1:
+            return None, {
+                "success": False,
+                "error": {"code": -1, "message": "Export task failed on NAS"},
+            }
+
+        time.sleep(poll_interval_sec)
+
+    return None, {
+        "success": False,
+        "error": {"code": -2, "message": "Timed out waiting for export"},
+    }
+
+
+def download_export(sid: str, dlid: int, event_name: str, file_ext: str, output_dir: str):
+    """
+    Downloads completed export data and saves to output_dir.
+    """
+    url = f"http://{NAS_IP}:5000/webapi/entry.cgi"
+    params = {
+        "api": "SYNO.SurveillanceStation.Recording",
+        "method": "OnRangeExportDone",
+        "version": "6",
+        "_sid": sid,
+        "dlid": dlid,
+        "fileName": event_name,
+    }
+    r = requests.get(url, params=params, stream=True, timeout=120)
+    r.raise_for_status()
+
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = os.path.join(output_dir, f"{event_name}.{file_ext}")
+    with open(output_path, "wb") as f:
+        for chunk in r.iter_content(chunk_size=1024 * 1024):
+            if chunk:
+                f.write(chunk)
+    return output_path
+
+
+def upload_to_filestation(local_path: str, nas_folder: str) -> str:
+    """
+    Uploads a local file to a NAS folder via SYNO.FileStation.Upload.
+    Uses a FileStation-scoped session (separate from SurveillanceStation).
+    Returns the NAS destination path: nas_folder/filename.
+    """
+    fs_sid, fs_token = get_sid(session="FileStation", request_token=True)
+    try:
+        url = f"http://{NAS_IP}:5000/webapi/entry.cgi"
+        filename = os.path.basename(local_path)
+        headers = {}
+        if fs_token:
+            headers["X-SYNO-TOKEN"] = fs_token
+        with open(local_path, "rb") as f:
+            r = requests.post(
+                url,
+                params={
+                    "api": "SYNO.FileStation.Upload",
+                    "version": "2",
+                    "method": "upload",
+                    "_sid": fs_sid,
+                },
+                data={
+                    "path": nas_folder,
+                    "create_parents": "true",
+                    "overwrite": "true",
+                },
+                files={"file": (filename, f, "application/octet-stream")},
+                headers=headers,
+                timeout=300,
+            )
+        r.raise_for_status()
+        payload = r.json()
+        if not payload.get("success"):
+            raise RuntimeError(f"FileStation upload failed: {payload}")
+        return f"{nas_folder.rstrip('/')}/{filename}"
+    finally:
+        logout_sid(fs_sid)
+
+
+def build_public_clip_ref(local_path: str, output_dir: str) -> str:
+    """
+    Convert local clip path to a shareable link if CLIP_PUBLIC_BASE is configured.
+    Example:
+      output_dir=/volume1/SharedFolder/EventClips
+      CLIP_PUBLIC_BASE=smb://nas.local/SharedFolder/EventClips
+      -> smb://nas.local/SharedFolder/EventClips/<file>
+    """
+    if not CLIP_PUBLIC_BASE:
+        return local_path
+    try:
+        rel = os.path.relpath(local_path, output_dir)
+    except ValueError:
+        # Different drives or invalid relpath on some platforms -> fallback
+        rel = os.path.basename(local_path)
+    rel = rel.replace("\\", "/")
+    return f"{CLIP_PUBLIC_BASE.rstrip('/')}/{rel.lstrip('/')}"
+
+
+def list_cameras(sid: str):
+    url = f"http://{NAS_IP}:5000/webapi/entry.cgi"
+    params = {
+        "api": "SYNO.SurveillanceStation.Camera",
+        "method": "List",
+        "version": "9",
+        "_sid": sid,
+        "limit": -1,
+    }
+    r = requests.get(url, params=params, timeout=30)
+    r.raise_for_status()
+    return r.json()
+
+
+def parse_time_to_unix(value: str) -> int:
+    """Accepts 'YYYY-MM-DD HH:MM:SS' or ISO 8601 and returns Unix timestamp."""
+    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"):
+        try:
+            return int(datetime.datetime.strptime(value, fmt).timestamp())
+        except ValueError:
+            pass
+    try:
+        iso_dt = datetime.datetime.fromisoformat(value.replace("Z", "+00:00"))
+        return int(iso_dt.timestamp())
+    except ValueError as exc:
+        raise ValueError(f"Invalid datetime format: {value}") from exc
+
+
+def run_single_export(
+    start_str: str,
+    end_str: str,
+    camera_id: int,
+    event_name: str,
+    output_dir: str,
+    nas_save: bool = False,
+    nas_upload: bool = False,
+    nas_dstdir: str = NAS_EXPORT_DSTDIR,
+):
+    start_unix = parse_time_to_unix(start_str)
+    end_unix = parse_time_to_unix(end_str)
+    if end_unix <= start_unix:
+        raise ValueError("End time must be later than start time.")
+
+    session_id = None
+    try:
+        session_id, _ = get_sid()
+        print("Logged in.")
+        if nas_save:
+            print(f"Creating NAS-side clip export in dstdir='{nas_dstdir}'...")
+            result, check_payload = start_nas_save_export(
+                session_id,
+                camera_id=camera_id,
+                start_time=start_unix,
+                end_time=end_unix,
+                event_name=event_name,
+                nas_dstdir=nas_dstdir,
+            )
+            if not result.get("success"):
+                check_result = (check_payload or {}).get("data", {}).get("result")
+                if check_result == 2:
+                    raise RuntimeError(
+                        "Failed to start NAS save export: no Surveillance Station event in selected time range "
+                        "(CheckAvailableExport result=2). "
+                        "Save-mode exports events only on this NAS build. "
+                        "Use export mode without --nas-save for arbitrary time ranges."
+                    )
+                raise RuntimeError(f"Failed to start NAS save export: {result}")
+            nas_ref = f"{nas_dstdir.rstrip('/')}/{event_name}"
+            print(f"NAS export task started successfully: {result}")
+            if CLIP_PUBLIC_BASE:
+                print(f"Share link prefix: {CLIP_PUBLIC_BASE.rstrip('/')}/{event_name}")
+            return nas_ref
+        else:
+            print("Creating clip...")
+            result = start_range_export(
+                session_id,
+                camera_id=camera_id,
+                start_time=start_unix,
+                end_time=end_unix,
+                event_name=event_name,
+            )
+
+            if not result.get("success"):
+                raise RuntimeError(f"Failed to start export: {result}")
+
+            dlid = result.get("data", {}).get("dlid")
+            if dlid is None:
+                raise RuntimeError(f"Export started but no dlid returned: {result}")
+
+            file_ext, progress_result = wait_for_export(session_id, dlid)
+            if not file_ext:
+                raise RuntimeError(f"Error while waiting for export: {progress_result}")
+
+            if nas_upload:
+                os.makedirs(CLIP_TMP_DIR, exist_ok=True)
+                tmp_path = os.path.join(CLIP_TMP_DIR, f"{event_name}.{file_ext}")
+                try:
+                    download_export(session_id, dlid, event_name, file_ext, CLIP_TMP_DIR)
+                    # Log out SS session before FileStation login to avoid duplicate-session conflicts
+                    logout_sid(session_id)
+                    session_id = None
+                    print(f"Downloaded to temp: {tmp_path}, uploading to NAS:{nas_dstdir} ...")
+                    nas_ref = upload_to_filestation(tmp_path, nas_dstdir)
+                    print(f"Success! Clip uploaded to NAS: {nas_ref}")
+                    if CLIP_PUBLIC_BASE:
+                        print(f"Share link: {CLIP_PUBLIC_BASE.rstrip('/')}/{os.path.basename(nas_ref)}")
+                    return nas_ref
+                finally:
+                    try:
+                        os.remove(tmp_path)
+                    except Exception:
+                        pass
+            else:
+                output_path = download_export(
+                    session_id,
+                    dlid,
+                    event_name,
+                    file_ext,
+                    output_dir,
+                )
+                public_ref = build_public_clip_ref(output_path, output_dir)
+                print(f"Success! Clip saved to {output_path}")
+                if public_ref != output_path:
+                    print(f"Share link: {public_ref}")
+                return output_path
+    finally:
+        if session_id:
+            logout_sid(session_id)
+
+
+def update_remote_status(job_id: int, status: str, output_path: str | None = None, error_message: str | None = None):
+    payload: dict[str, Any] = {
+        "job_id": job_id,
+        "status": status,
+    }
+    if output_path is not None:
+        payload["output_path"] = output_path
+    if error_message is not None:
+        payload["error_message"] = error_message
+
+    headers = {
+        "Content-Type": "application/json",
+        "X-Shared-Password": AAREDB_SHARED_PASSWORD or "",
+    }
+    r = requests.post(
+        DISPATCHER_STATUS_URL,
+        data=json.dumps(payload),
+        headers=headers,
+        verify=DISPATCHER_VERIFY_SSL,
+        cert=(DISPATCHER_CLIENT_CERT, DISPATCHER_CLIENT_KEY) if DISPATCHER_CLIENT_CERT else None,
+        timeout=60,
+    )
+    r.raise_for_status()
+    return r.json()
+
+
+def build_event_name(job: dict[str, Any]) -> str:
+    sample_id = job.get("sample_id", "s")
+    camera_id = job.get("camera_id", "c")
+    job_id = job.get("job_id", "j")
+    start_time = str(job.get("start_time", "")).replace(":", "-").replace(" ", "_")
+    return f"clip_sample{sample_id}_cam{camera_id}_job{job_id}_{start_time}"
+
+
+def process_job(job: dict[str, Any], output_dir: str, nas_upload: bool = False, nas_dstdir: str = NAS_EXPORT_DSTDIR):
+    job_id = int(job["job_id"])
+    camera_id = int(job["camera_id"])
+    start_time = str(job["start_time"])
+    end_time = str(job["end_time"])
+    event_name = build_event_name(job)
+
+    print(f"[JOB {job_id}] Claiming job (camera={camera_id}, start={start_time}, end={end_time})")
+    update_remote_status(job_id, "submitted")
+
+    try:
+        output_path = run_single_export(
+            start_str=start_time,
+            end_str=end_time,
+            camera_id=camera_id,
+            event_name=event_name,
+            output_dir=output_dir,
+            nas_upload=nas_upload,
+            nas_dstdir=nas_dstdir,
+        )
+        public_ref = build_public_clip_ref(output_path, output_dir) if not nas_upload else output_path
+        print(f"[JOB {job_id}] Done -> {output_path}")
+        update_remote_status(job_id, "done", output_path=public_ref, error_message=None)
+    except Exception as exc:
+        err = str(exc)
+        print(f"[JOB {job_id}] Failed: {err}")
+        update_remote_status(job_id, "failed", error_message=err)
+
+
+def worker_loop(output_dir: str, nas_upload: bool = False, nas_dstdir: str = NAS_EXPORT_DSTDIR):
+    print(f"Connecting to clip jobs websocket: {DISPATCHER_WS_URL}")
+    active_jobs: set[int] = set()
+
+    while True:
+        ws = None
+        try:
+            ws = create_connection(
+                DISPATCHER_WS_URL,
+                header=[f"X-Shared-Password: {AAREDB_SHARED_PASSWORD or ''}"],
+                timeout=30,
+                sslopt=(
+                    {
+                        "ca_certs": DISPATCHER_CA_CERT,
+                        "certfile": DISPATCHER_CLIENT_CERT,
+                        "keyfile": DISPATCHER_CLIENT_KEY,
+                        "cert_reqs": 2,  # ssl.CERT_REQUIRED
+                    }
+                    if DISPATCHER_CA_CERT
+                    else (None if DISPATCHER_VERIFY_SSL else {"cert_reqs": 0})
+                ),
+            )
+            print("Websocket connected. Waiting for clip jobs...")
+            while True:
+                raw = ws.recv()
+                if not raw:
+                    continue
+                try:
+                    payload = json.loads(raw)
+                except json.JSONDecodeError:
+                    print(f"Skipping non-JSON websocket payload: {raw!r}")
+                    continue
+
+                if not isinstance(payload, list):
+                    continue
+
+                to_do_jobs = [
+                    j
+                    for j in payload
+                    if isinstance(j, dict)
+                    and str(j.get("status", "")).lower() == "to_do"
+                    and j.get("job_id") is not None
+                ]
+
+                for job in sorted(to_do_jobs, key=lambda x: int(x.get("job_id", 0))):
+                    job_id = int(job["job_id"])
+                    if job_id in active_jobs:
+                        continue
+                    active_jobs.add(job_id)
+                    try:
+                        process_job(job, output_dir=output_dir, nas_upload=nas_upload, nas_dstdir=nas_dstdir)
+                    finally:
+                        active_jobs.discard(job_id)
+        except Exception as exc:
+            print(f"Worker websocket error: {exc}. Reconnecting in 5 seconds...")
+            time.sleep(5)
+        finally:
+            try:
+                if ws:
+                    ws.close()
+            except Exception:
+                pass
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Surveillance Station clip tool: export once or run as websocket worker."
+    )
+    sub = parser.add_subparsers(dest="mode")
+
+    export_parser = sub.add_parser("export", help="Export one clip from explicit start/end")
+    export_parser.add_argument("start", help='Start time "YYYY-MM-DD HH:MM:SS" or ISO8601')
+    export_parser.add_argument("end", help='End time "YYYY-MM-DD HH:MM:SS" or ISO8601')
+    export_parser.add_argument("--camera-id", type=int, default=1)
+    export_parser.add_argument("--event-name", default="Event_101")
+    export_parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR)
+    export_parser.add_argument("--list-cameras", action="store_true")
+    export_parser.add_argument(
+        "--nas-save",
+        action="store_true",
+        help="Ask NAS to save clip locally on NAS instead of downloading to this machine",
+    )
+    export_parser.add_argument(
+        "--nas-dstdir",
+        default=NAS_EXPORT_DSTDIR,
+        help="Destination directory for NAS-side save export (default from NAS_EXPORT_DSTDIR/EventClips)",
+    )
+
+    worker_parser = sub.add_parser("worker", help="Connect websocket and process clip jobs")
+    worker_parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR)
+    worker_parser.add_argument(
+        "--nas-upload",
+        action="store_true",
+        help="Download clip to temp then upload to NAS via FileStation (works for any time range)",
+    )
+    worker_parser.add_argument(
+        "--nas-dstdir",
+        default=NAS_EXPORT_DSTDIR,
+        help="Destination folder on NAS for --nas-upload",
+    )
+
+    # Backward compatibility with previous invocation style:
+    # python bl-clip-manager.py "<start>" "<end>"
+    parser.add_argument("legacy_start", nargs="?", help=argparse.SUPPRESS)
+    parser.add_argument("legacy_end", nargs="?", help=argparse.SUPPRESS)
+    parser.add_argument("--camera-id", type=int, default=1, help=argparse.SUPPRESS)
+    parser.add_argument("--event-name", default="Event_101", help=argparse.SUPPRESS)
+    parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR, help=argparse.SUPPRESS)
+    parser.add_argument("--list-cameras", action="store_true", help=argparse.SUPPRESS)
+    parser.add_argument("--nas-save", action="store_true", help=argparse.SUPPRESS)
+    parser.add_argument("--nas-dstdir", default=NAS_EXPORT_DSTDIR, help=argparse.SUPPRESS)
+
+    return parser.parse_args()
+
+
+def main():
+    ensure_env_ready()
+    args = parse_args()
+
+    # Legacy mode support
+    if args.mode is None and args.legacy_start and args.legacy_end:
+        if args.list_cameras:
+            sid, _ = get_sid()
+            try:
+                camera_result = list_cameras(sid)
+                if camera_result.get("success"):
+                    cams = camera_result.get("data", {}).get("cameras", [])
+                    print("Available cameras:")
+                    for cam in cams:
+                        print(f"- id={cam.get('id')} name={cam.get('name')}")
+                else:
+                    print(f"Could not list cameras: {camera_result}")
+            finally:
+                logout_sid(sid)
+            return
+        run_single_export(
+            start_str=args.legacy_start,
+            end_str=args.legacy_end,
+            camera_id=args.camera_id,
+            event_name=args.event_name,
+            output_dir=args.output_dir,
+            nas_save=args.nas_save,
+            nas_dstdir=args.nas_dstdir,
+        )
+        return
+
+    if args.mode == "export":
+        if args.list_cameras:
+            sid, _ = get_sid()
+            try:
+                camera_result = list_cameras(sid)
+                if camera_result.get("success"):
+                    cams = camera_result.get("data", {}).get("cameras", [])
+                    print("Available cameras:")
+                    for cam in cams:
+                        print(f"- id={cam.get('id')} name={cam.get('name')}")
+                else:
+                    print(f"Could not list cameras: {camera_result}")
+            finally:
+                logout_sid(sid)
+            return
+
+        run_single_export(
+            start_str=args.start,
+            end_str=args.end,
+            camera_id=args.camera_id,
+            event_name=args.event_name,
+            output_dir=args.output_dir,
+            nas_save=args.nas_save,
+            nas_dstdir=args.nas_dstdir,
+        )
+        return
+
+    if args.mode == "worker":
+        worker_loop(output_dir=args.output_dir, nas_upload=args.nas_upload, nas_dstdir=args.nas_dstdir)
+        return
+
+    raise SystemExit(
+        "Usage:\n"
+        "  python bl-clip-manager.py worker\n"
+        "  python bl-clip-manager.py export \"YYYY-MM-DD HH:MM:SS\" \"YYYY-MM-DD HH:MM:SS\"\n"
+        "  python bl-clip-manager.py \"YYYY-MM-DD HH:MM:SS\" \"YYYY-MM-DD HH:MM:SS\"  (legacy)"
+    )
+
+
+if __name__ == "__main__":
+    main()

src/core/__init__.py

@@ -0,0 +1,12 @@
+"""Core functionality for AareLC ML Studio."""
+from .zmq_client import ZMQStreamClient
+from .inference_client import InferenceClient
+from .db_client import DatabaseClient
+from .image_processor import ImageProcessor
+
+__all__ = [
+    'ZMQStreamClient',
+    'InferenceClient',
+    'DatabaseClient',
+    'ImageProcessor'
+]

src/core/cv_logic.py

@@ -0,0 +1,71 @@
+"""Computer vision logic utilities for edge detection and polygon refinement.
+
+This module is maintained for backward compatibility.
+For new code, please use src.core.image_processor.ImageProcessor instead.
+"""
+import cv2
+import numpy as np
+
+
+def get_refined_polygon(self, roi, p):
+    """
+    Robust contour detection optimized for homogeneous backgrounds.
+    Uses Adaptive Thresholding + Hole Filling.
+    """
+    # 1. Preprocessing
+    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
+    
+    # Median blur is superior for preserving edges on homogeneous backgrounds
+    blurred = cv2.medianBlur(gray, 5)
+    
+    # 2. Adaptive Thresholding
+    # block_size must be odd and > 1. We use morph_kernel slider to control this.
+    block_size = max(3, p["morph"] if p["morph"] % 2 != 0 else p["morph"] + 1)
+    
+    # Gaussian adaptive thresholding is very resilient to homogeneous background noise
+    thresh = cv2.adaptiveThreshold(
+        blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+        cv2.THRESH_BINARY_INV, block_size, 2
+    )
+
+    # 3. Clean up noise
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    opened = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
+        
+    # 4. Fill Holes (Crucial for loops)
+    # We find all contours and draw them filled on a mask
+    cnts, _ = cv2.findContours(opened, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
+    mask = np.zeros_like(opened)
+    for c in cnts:
+        cv2.drawContours(mask, [c], -1, 255, -1) # -1 thickness = fill
+
+    # 5. Final Contour Extraction
+    # Now we take the largest filled blob
+    final_cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+    if not final_cnts:
+        return np.array([]), opened
+
+    # Select largest blob by area
+    cnt = max(final_cnts, key=cv2.contourArea)
+        
+    # 6. Smooth the polygon based on the user's precision slider
+    epsilon = p["eps"] * cv2.arcLength(cnt, True)
+    approx = cv2.approxPolyDP(cnt, epsilon, True)
+        
+    return approx, mask
+
+
+def measure_focus(raw_frame):
+    """
+    Measure focus quality using Laplacian variance.
+
+    DEPRECATED: Use ImageProcessor.measure_focus() instead.
+
+    Args:
+        raw_frame: Grayscale image
+
+    Returns:
+        Focus score (higher = better focus)
+    """
+    return cv2.Laplacian(raw_frame, cv2.CV_64F).var()

src/core/db_client.py

@@ -0,0 +1,296 @@
+"""Client for database server communication."""
+import requests
+import numpy as np
+import cv2
+from pathlib import Path
+
+
+class DatabaseClient:
+    """Handles communication with database server for image retrieval."""
+
+    def __init__(self, server_url, shared_password="", download_base=None):
+        """
+        Initialize database client.
+
+        Args:
+            server_url: Base URL of database server (e.g., "http://localhost:8000")
+            shared_password: Password for shared auth
+            download_base: Base URL for image downloads (e.g., "http://localhost:8002/backend")
+        """
+        self.server_url = server_url
+        self.shared_password = shared_password
+        self.download_base = download_base or "http://localhost:8002/backend"
+
+    def _get_headers(self):
+        return {"x-shared-password": self.shared_password}
+
+    def fetch_next_raw_image(self, timeout=5):
+        # 1. Get metadata from Dispatcher (via Tunnel 8001 or Direct 443)
+        # Assuming you've set self.server_url to the dispatcher base
+        meta_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/raw_next"
+        r = requests.get(meta_url, headers=self._get_headers(), verify=False, timeout=timeout)
+        
+        if r.status_code != 200:
+            raise ValueError(f"No more images available (status {r.status_code})")
+
+        data = r.json()
+        image_obj = data.get('image', {})
+        
+        # 2. Construct the correct download URL
+        filepath = image_obj.get('filepath')
+        if not filepath:
+            raise ValueError("Metadata received but no filepath found.")
+
+        # Use the configurable download_base
+        download_url = f"{self.download_base.rstrip('/')}/{filepath.lstrip('/')}"
+        
+        print(f"DEBUG: Fetching pixels from: {download_url}")
+        img_res = requests.get(download_url, timeout=timeout, verify=False)
+
+        if img_res.status_code != 200:
+            # Fallback debug
+            raise ValueError(f"Could not retrieve image from {download_url} (status {img_res.status_code})")
+
+        # 3. Process the content
+        content = img_res.content
+
+        # Use OpenCV to decode
+        nparr = np.frombuffer(content, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        if img is None:
+            # If standard decode fails, check if it's a RAW Bayer file (which your system uses)
+            # A 4MP Bayer file is exactly 4194304 bytes or 4177936 bytes
+            if len(content) in [4194304, 4177936]:
+                try:
+                    # Attempt raw Bayer recovery (2048x2048 or 2048x2040)
+                    h = 2048
+                    w = len(content) // h
+                    raw = np.frombuffer(content, np.uint8).reshape((h, w))
+                    img = cv2.cvtColor(raw, cv2.COLOR_BAYER_RG2BGR)
+                    print(f"DEBUG: Successfully decoded as raw Bayer {w}x{h}")
+                except:
+                    pass
+
+        if img is None:
+            raise ValueError(f"Failed to decode image data. size={len(content)} bytes. Header: {content[:4].hex()}")
+
+        # Return both image and metadata
+        return img, data
+
+    def set_server_url(self, url):
+        """Update server URL."""
+        self.server_url = url
+
+    def get_server_url(self):
+        """Get current server URL."""
+        return self.server_url
+
+    def set_shared_password(self, password):
+        """Update the shared password."""
+        self.shared_password = password
+
+    def set_download_base(self, url):
+        """Update the image download base URL."""
+        self.download_base = url
+
+    def save_annotation(self, image_id, segments, username=None, timeout=10):
+        """
+        Save YOLO segmentation annotations to database.
+
+        Args:
+            image_id: ID of the image in the database
+            segments: List of dicts with 'class_id' and 'points' (normalized YOLO format)
+            username: Optional username to attribute the annotation to
+            timeout: Request timeout in seconds
+
+        Returns:
+            Response dict with status and annotation_id
+        """
+        save_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/{image_id}"
+
+        payload = {
+            "segments": segments,
+            "username": username
+        }
+
+        print(f"DEBUG: Saving annotation to: {save_url}")
+        print(f"DEBUG: Payload: {len(segments)} segments")
+
+        response = requests.post(
+            save_url,
+            json=payload,
+            headers=self._get_headers(),
+            verify=False,
+            timeout=timeout
+        )
+
+        if response.status_code != 200:
+            raise ValueError(f"Failed to save annotation: {response.status_code} - {response.text}")
+
+        return response.json()
+
+    def skip_image(self, image_id, timeout=5):
+        """
+        Mark an image as skipped in the database.
+
+        Args:
+            image_id: ID of the image to skip
+            timeout: Request timeout in seconds
+
+        Returns:
+            Response dict with status
+        """
+        skip_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/{image_id}/skip"
+
+        print(f"DEBUG: Skipping image: {skip_url}")
+
+        response = requests.post(
+            skip_url,
+            headers=self._get_headers(),
+            verify=False,
+            timeout=timeout
+        )
+
+        if response.status_code != 200:
+            raise ValueError(f"Failed to skip image: {response.status_code} - {response.text}")
+
+        return response.json()
+
+
+    def sync_local_annotations(
+        self,
+        since=None,
+        until=None,
+        include_unannotated=False,
+        dry_run=False,
+        limit=5000,
+        timeout=120,
+    ):
+        """Retrieve DB sync items and write missing files locally on this machine."""
+        sync_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/sync/local"
+
+        params = {
+            "include_unannotated": str(bool(include_unannotated)).lower(),
+            "limit": int(limit),
+        }
+        if since:
+            params["since"] = since
+        if until:
+            params["until"] = until
+
+        print(f"DEBUG: Syncing local files from: {sync_url}")
+        print(f"DEBUG: Sync params: {params}")
+
+        response = requests.post(
+            sync_url,
+            params=params,
+            headers=self._get_headers(),
+            verify=False,
+            timeout=timeout,
+        )
+
+        if response.status_code != 200:
+            raise ValueError(
+                f"Failed to sync local files: {response.status_code} - {response.text}"
+            )
+        data = response.json()
+        items = data.get("items", [])
+
+        images_dir = Path("data") / "img_sets" / "seg_training" / "images"
+        labels_dir = Path("data") / "img_sets" / "seg_training" / "labels"
+        if not dry_run:
+            images_dir.mkdir(parents=True, exist_ok=True)
+            labels_dir.mkdir(parents=True, exist_ok=True)
+
+        stats = {
+            "scanned": len(items),
+            "image_copied": 0,
+            "image_existing": 0,
+            "image_failed": 0,
+            "image_skipped_no_label": 0,
+            "label_copied": 0,
+            "label_existing": 0,
+            "label_failed": 0,
+            "dry_run": bool(dry_run),
+            "since": since,
+            "until": until,
+        }
+        failures = []
+
+        for item in items:
+            image_id = item.get("image_id")
+            image_url = item.get("image_url")
+            segments = item.get("segments") or []
+
+            if image_id is None:
+                stats["image_failed"] += 1
+                failures.append("missing image_id in sync item")
+                continue
+
+            image_path = images_dir / f"img_{int(image_id)}.jpg"
+            label_path = labels_dir / f"img_{int(image_id)}.txt"
+
+            lines = []
+            for seg in segments:
+                points = seg.get("points")
+                if (
+                    not isinstance(points, list)
+                    or len(points) < 6
+                    or len(points) % 2 != 0
+                ):
+                    continue
+                class_id = int(seg.get("class_id", 0))
+                lines.append(f"{class_id} " + " ".join(f"{float(p):.6f}" for p in points))
+
+            label_ready = False
+            if label_path.exists():
+                stats["label_existing"] += 1
+                label_ready = True
+            elif not lines:
+                stats["label_failed"] += 1
+                stats["image_skipped_no_label"] += 1
+                failures.append(
+                    f"image_id={image_id}: no valid segmentation points for label"
+                )
+                continue
+            else:
+                try:
+                    if not dry_run:
+                        label_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
+                    stats["label_copied"] += 1
+                    label_ready = True
+                except Exception as exc:
+                    stats["label_failed"] += 1
+                    failures.append(
+                        f"image_id={image_id}: failed writing {label_path} ({exc})"
+                    )
+                    continue
+
+            if not label_ready:
+                stats["image_skipped_no_label"] += 1
+                continue
+
+            if image_path.exists():
+                stats["image_existing"] += 1
+            else:
+                try:
+                    img_res = requests.get(
+                        image_url,
+                        timeout=min(timeout, 30),
+                        verify=False,
+                    )
+                    if img_res.status_code != 200 or not img_res.content:
+                        raise ValueError(
+                            f"image download failed status={img_res.status_code}"
+                        )
+                    if not dry_run:
+                        image_path.write_bytes(img_res.content)
+                    stats["image_copied"] += 1
+                except Exception as exc:
+                    stats["image_failed"] += 1
+                    failures.append(
+                        f"image_id={image_id}: failed downloading {image_url} ({exc})"
+                    )
+
+        return {"status": "ok", "stats": stats, "failures": failures[:100]}

src/core/image_processor.py

@@ -0,0 +1,586 @@
+import cv2
+import numpy as np
+
+
+class ImageProcessor:
+    def __init__(self):
+        self.edge_map = None
+        self.edge_map_shape = None
+        self.roi_offset = (0, 0) # Store (x1, y1) of the box
+
+    def compute_edge_map(self, image, box=None):
+        """
+        Compute edge map, optionally restricted to a specific box.
+        """
+        if image is None:
+            return None
+
+        # 1. Handle ROI crop if box is provided
+        if box:
+            x1, y1, x2, y2 = int(box['x1']), int(box['y1']), int(box['x2']), int(box['y2'])
+            # Add small padding to capture edges right at the boundary
+            h, w = image.shape[:2]
+            p = 10
+            x1, y1 = max(0, x1-p), max(0, y1-p)
+            x2, y2 = min(w, x2+p), min(h, y2+p)
+            
+            working_img = image[y1:y2, x1:x2]
+            self.roi_offset = (x1, y1)
+        else:
+            working_img = image
+            self.roi_offset = (0, 0)
+
+        # 2. Convert and process local ROI
+        gray = cv2.cvtColor(working_img, cv2.COLOR_BGR2GRAY)
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+
+        adaptive = cv2.adaptiveThreshold(
+            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
+            cv2.THRESH_BINARY_INV, 11, 2
+        )
+        canny = cv2.Canny(blurred, 30, 100)
+        edge_map = cv2.bitwise_or(adaptive, canny)
+        
+        kernel = np.ones((3,3), np.uint8)
+        self.edge_map = cv2.dilate(edge_map, kernel, iterations=1)
+        self.edge_map_shape = working_img.shape[:2]
+        return self.edge_map
+
+    def find_nearest_edge(self, x, y, radius=20):
+        """
+        Find nearest edge within the stored ROI.
+        """
+        if self.edge_map is None:
+            return x, y
+
+        # Translate global mouse coords to ROI local coords
+        lx, ly = int(x - self.roi_offset[0]), int(y - self.roi_offset[1])
+        h, w = self.edge_map.shape
+
+        # Bounds check
+        x0, x1 = max(0, lx - radius), min(w - 1, lx + radius)
+        y0, y1 = max(0, ly - radius), min(h - 1, ly + radius)
+
+        region = self.edge_map[y0:y1+1, x0:x1+1]
+        edge_points = np.argwhere(region > 128)
+
+        if len(edge_points) == 0:
+            return x, y
+
+        # Find closest point in local ROI space
+        edge_points_local = edge_points + [y0, x0]
+        distances = np.linalg.norm(edge_points_local - [ly, lx], axis=1)
+        closest_idx = np.argmin(distances)
+        closest_y, closest_x = edge_points_local[closest_idx]
+
+        # Translate back to global image coords
+        return int(closest_x + self.roi_offset[0]), int(closest_y + self.roi_offset[1])
+
+    @staticmethod
+    def get_refined_polygon_watershed(roi, p, manual_mask=None):
+        """
+        Watershed-based segmentation using manual strokes as markers.
+        More intuitive than GrabCut: background strokes act as barriers.
+        """
+        h, w = roi.shape[:2]
+        print(f"  Watershed: roi_shape={roi.shape}")
+        if h < 10 or w < 10:
+            print(f"  Watershed: ROI too small ({h}x{w})")
+            return np.array([]), np.zeros((h, w), dtype=np.uint8)
+
+        # Convert to grayscale for processing
+        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
+
+        # Compute gradient magnitude (edges)
+        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
+        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
+        gradient = np.sqrt(sobelx**2 + sobely**2)
+        gradient = cv2.normalize(gradient, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+
+        # Create markers for watershed
+        markers = np.zeros((h, w), dtype=np.int32)
+
+        if manual_mask is not None:
+            # Match sizes
+            if manual_mask.shape[:2] != (h, w):
+                print(f"  Watershed: Resizing manual_mask from {manual_mask.shape[:2]} to ({h}, {w})")
+                manual_mask = cv2.resize(manual_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+
+            has_bg = np.any(manual_mask == 0)
+            has_fg = np.any(manual_mask == 1)
+            print(f"  Watershed: has_bg={has_bg}, has_fg={has_fg}")
+
+            if has_bg or has_fg:
+                # Background strokes = marker 1 (will be excluded)
+                if has_bg:
+                    markers[manual_mask == 0] = 1
+
+                # Foreground strokes = marker 2 (will be included)
+                if has_fg:
+                    fg_pixels = np.sum(manual_mask == 1)
+                    print(f"  Watershed: Using {fg_pixels} foreground pixels")
+                    markers[manual_mask == 1] = 2
+
+                    # If only foreground painted, automatically add background border
+                    if not has_bg:
+                        print(f"  Watershed: No background painted - using border as background hint")
+                        border_width = max(3, min(w, h) // 30)
+                        markers[:border_width, :] = 1
+                        markers[-border_width:, :] = 1
+                        markers[:, :border_width] = 1
+                        markers[:, -border_width:] = 1
+                else:
+                    # No foreground marked - use center region
+                    center_x, center_y = w // 2, h // 2
+                    radius = min(w, h) // 6
+                    y_indices, x_indices = np.ogrid[:h, :w]
+                    center_mask = (x_indices - center_x)**2 + (y_indices - center_y)**2 <= radius**2
+                    markers[center_mask] = 2
+                    print(f"  Watershed: Using center circle as foreground hint")
+
+                # Apply watershed
+                print(f"  Watershed: Running watershed algorithm...")
+                markers = cv2.watershed(roi, markers)
+                print(f"  Watershed: Watershed complete, unique markers: {np.unique(markers)}")
+
+                # Extract foreground (marker == 2)
+                binary_mask = np.where(markers == 2, 255, 0).astype(np.uint8)
+                fg_area = np.sum(binary_mask > 0)
+                print(f"  Watershed: Binary mask has {fg_area} foreground pixels")
+            else:
+                # No manual strokes - use automatic thresholding
+                binary_mask = ImageProcessor._get_auto_mask(roi, p)
+        else:
+            # No manual mask - use automatic thresholding
+            binary_mask = ImageProcessor._get_auto_mask(roi, p)
+
+        # Morphological cleanup
+        k = max(1, p.get("morph", 7))
+        k = k if k % 2 != 0 else k + 1
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
+        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
+        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel)
+
+        # Find contours
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        print(f"  Watershed: Found {len(contours)} contours")
+        if not contours:
+            print(f"  Watershed: No contours found!")
+            return np.array([]), binary_mask
+
+        # Get largest contour
+        cnt = max(contours, key=cv2.contourArea)
+        area = cv2.contourArea(cnt)
+        print(f"  Watershed: Largest contour area: {area}")
+
+        # Smooth polygon
+        epsilon = p.get("eps", 0.01) * cv2.arcLength(cnt, True)
+        approx = cv2.approxPolyDP(cnt, epsilon, True)
+        print(f"  Watershed: Polygon has {len(approx)} points")
+
+        return approx, binary_mask
+
+    @staticmethod
+    def _get_auto_mask(roi, p):
+        """Automatic mask generation without manual strokes."""
+        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+
+        # Adaptive thresholding
+        adaptive = cv2.adaptiveThreshold(
+            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+            cv2.THRESH_BINARY_INV, 11, 2
+        )
+
+        # Canny edges
+        canny = cv2.Canny(blurred, 30, p.get("high", 150))
+
+        # Combine
+        binary_mask = cv2.bitwise_or(adaptive, canny)
+
+        return binary_mask
+
+    @staticmethod
+    def get_refined_polygon(roi, p, manual_mask=None):
+        """
+        Refinement using Interactive GrabCut.
+        Optimized for homogeneous backgrounds by using the ROI border as a background seed.
+        """
+        h, w = roi.shape[:2]
+        if h < 10 or w < 10: return np.array([]), np.zeros((h, w), dtype=np.uint8)
+
+        pad_w, pad_h = int(w * 0.05), int(h * 0.05)
+        rect = (pad_w, pad_h, w - 2 * pad_w, h - 2 * pad_h)
+
+        # Initialize mask with proper GrabCut values
+        if manual_mask is not None:
+            # Match sizes to prevent IndexErrors
+            if manual_mask.shape[:2] != roi.shape[:2]:
+                manual_mask = cv2.resize(manual_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+
+            # Check what strokes we have
+            has_bg = np.any(manual_mask == 0)
+            has_fg = np.any(manual_mask == 1)
+
+            # GrabCut needs BOTH foreground and background samples
+            if has_bg and not has_fg:
+                # User only painted background - use rectangle center as foreground hint
+                print("  GrabCut mode: Background only - adding center as foreground hint")
+                mask = np.full((h, w), cv2.GC_PR_BGD, dtype=np.uint8)  # Start with probable background
+                mask[manual_mask == 0] = cv2.GC_BGD  # Hard background from user
+                # Set center region as probable foreground
+                center_x, center_y = w // 2, h // 2
+                radius = min(w, h) // 4
+                y_indices, x_indices = np.ogrid[:h, :w]
+                center_mask = (x_indices - center_x)**2 + (y_indices - center_y)**2 <= radius**2
+                mask[center_mask] = cv2.GC_PR_FGD
+                print(f"    BG pixels: {np.sum(mask == cv2.GC_BGD)}, FG hint pixels: {np.sum(mask == cv2.GC_PR_FGD)}")
+            elif has_fg and not has_bg:
+                # User only painted foreground - use border as background hint
+                mask = np.full((h, w), cv2.GC_PR_FGD, dtype=np.uint8)  # Start with probable foreground
+                mask[manual_mask == 1] = cv2.GC_FGD  # Hard foreground from user
+                # Set border as probable background
+                border_width = max(2, min(w, h) // 20)
+                mask[:border_width, :] = cv2.GC_PR_BGD
+                mask[-border_width:, :] = cv2.GC_PR_BGD
+                mask[:, :border_width] = cv2.GC_PR_BGD
+                mask[:, -border_width:] = cv2.GC_PR_BGD
+            elif has_bg and has_fg:
+                # User painted both - use their constraints directly
+                mask = np.full((h, w), cv2.GC_PR_FGD, dtype=np.uint8)
+                mask[manual_mask == 0] = cv2.GC_BGD
+                mask[manual_mask == 1] = cv2.GC_FGD
+            else:
+                # No strokes at all - fall back to rectangle mode
+                mask = np.zeros((h, w), np.uint8)
+                mode = cv2.GC_INIT_WITH_RECT
+                iters = max(1, min(2, int(p.get("high", 150) / 75)))
+                bgdModel = np.zeros((1, 65), np.float64)
+                fgdModel = np.zeros((1, 65), np.float64)
+                try:
+                    cv2.grabCut(roi, mask, rect, bgdModel, fgdModel, iters, mode)
+                except Exception as e:
+                    print(f"GrabCut failed: {e}")
+                    return np.array([]), np.zeros((h, w), dtype=np.uint8)
+                binary_mask = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8') * 255
+                k = max(1, p.get("morph", 7))
+                k = k if k % 2 != 0 else k + 1
+                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
+                binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
+                contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                if not contours:
+                    return np.array([]), binary_mask
+                cnt = max(contours, key=cv2.contourArea)
+                epsilon = p.get("eps", 0.01) * cv2.arcLength(cnt, True)
+                approx = cv2.approxPolyDP(cnt, epsilon, True)
+                return approx, binary_mask
+
+            mode = cv2.GC_INIT_WITH_MASK
+            # More iterations when we have manual guidance
+            iters = 5
+        else:
+            mask = np.zeros((h, w), np.uint8)
+            mode = cv2.GC_INIT_WITH_RECT
+            iters = max(1, min(2, int(p.get("high", 150) / 75)))
+
+        bgdModel = np.zeros((1, 65), np.float64)
+        fgdModel = np.zeros((1, 65), np.float64)
+
+        try:
+            print(f"  Running GrabCut: mode={mode}, iters={iters}, roi_shape={roi.shape}")
+            cv2.grabCut(roi, mask, rect, bgdModel, fgdModel, iters, mode)
+            print(f"  GrabCut succeeded")
+        except Exception as e:
+            print(f"  GrabCut failed: {e}")
+            return np.array([]), np.zeros((h, w), dtype=np.uint8)
+
+        # Convert to binary mask
+        binary_mask = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8') * 255
+
+        # 4. Clean up the mask using the morphology slider
+        k = max(1, p.get("morph", 7))
+        k = k if k % 2 != 0 else k + 1
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
+        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
+
+        # 5. Find the largest contour
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if not contours:
+            return np.array([]), binary_mask
+
+        cnt = max(contours, key=cv2.contourArea)
+
+        # 6. Smooth the polygon based on the epsilon slider
+        epsilon = p.get("eps", 0.01) * cv2.arcLength(cnt, True)
+        approx = cv2.approxPolyDP(cnt, epsilon, True)
+
+        return approx, binary_mask
+
+    def draw_detections(self, image, detections, class_params, class_colors, active_class=None,
+                        show_debug_edges=False, refine_mode=False, selected_idx=-1, selected_island_idx=0,
+                        selected_island_kind="outer", hidden_classes=None,
+                        segmentation_method='watershed', annotation_line_width=2,
+                        show_boxes=True, show_segments=True):
+        """
+        Draw detection boxes and polygons.
+        Supports manual editing by checking for existing 'poly' in detections.
+        """
+        draw_img = image.copy()
+        hidden_classes = hidden_classes or set()
+
+        # Show global edge map overlay if requested
+        if show_debug_edges and self.edge_map is not None:
+            # Create colored edge overlay
+            edge_colored = cv2.cvtColor(self.edge_map, cv2.COLOR_GRAY2BGR)
+            edge_colored[:, :, 0] = 0  # Remove blue channel
+            edge_colored[:, :, 1] = 0  # Remove green channel
+            # Keep only red channel for edges
+            draw_img = cv2.addWeighted(draw_img, 0.7, edge_colored, 0.3, 0)
+
+        for i, det in enumerate(detections):
+            is_selected = (i == selected_idx)
+            try:
+                box = det.get('box', det)
+                x1, y1, x2, y2 = int(box.get('x1', 0)), int(box.get('y1', 0)), int(box.get('x2', 0)), int(
+                    box.get('y2', 0))
+                label = det.get('name', det.get('label', 'unknown'))
+            except:
+                continue
+
+            # Skip hidden classes
+            if label in hidden_classes:
+                continue
+
+            color = class_colors.get(label, (255, 0, 255))
+            params = class_params.get(label) or next(iter(class_params.values()), {})
+            base_thickness = max(1, int(annotation_line_width))
+            thickness = base_thickness + 1 if is_selected else base_thickness
+
+            # 1. Draw Bounding Box
+            if show_boxes:
+                cv2.rectangle(draw_img, (x1, y1), (x2, y2), color, thickness)
+                cv2.putText(draw_img, f"{label} {'(EDITING)' if is_selected else ''}",
+                            (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+
+            if is_selected and show_boxes:
+                # DRAW RESIZE HANDLES (Small white squares at all 4 corners)
+                handle_size = max(3, base_thickness + 2)
+                corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
+                for (cx, cy) in corners:
+                    cv2.rectangle(draw_img, (cx - handle_size, cy - handle_size),
+                                  (cx + handle_size, cy + handle_size), (255, 255, 255), -1)
+                    cv2.rectangle(draw_img, (cx - handle_size, cy - handle_size),
+                                  (cx + handle_size, cy + handle_size), color, max(1, base_thickness - 1))
+
+            # 2. Handle Polygon (Manual or Refined)
+            polygons = []
+            if isinstance(det.get('polygons'), list):
+                for poly_item in det.get('polygons'):
+                    arr = np.array(poly_item, dtype=np.float32).reshape(-1, 1, 2)
+                    if len(arr) >= 3:
+                        polygons.append(arr)
+            poly = det.get('poly')
+            if not polygons and poly is not None and len(poly) >= 3:
+                polygons = [np.array(poly, dtype=np.float32).reshape(-1, 1, 2)]
+            holes = []
+            if isinstance(det.get('holes'), list):
+                for hole_item in det.get('holes'):
+                    arr = np.array(hole_item, dtype=np.float32).reshape(-1, 1, 2)
+                    if len(arr) >= 3:
+                        holes.append(arr)
+            manual_mask = det.get('manual_mask') # Retrieve user brush strokes stored in the detection
+
+            # Regenerate polygon if: missing, empty, or needs refinement
+            suppress_auto_polygon = bool(det.get('suppress_auto_polygon', False))
+            should_regenerate = show_segments and refine_mode and y2 > y1 and x2 > x1 and (
+                len(polygons) == 0
+            ) and not suppress_auto_polygon
+
+            if should_regenerate:
+                print(f"  Should regenerate: method={segmentation_method}, manual_mask={'present' if manual_mask is not None else 'absent'}")
+                roi = image[y1:y2, x1:x2]
+                if roi.size > 0:
+                    # Use selected segmentation method
+                    if segmentation_method == 'watershed':
+                        generated_poly, debug_mask = self.get_refined_polygon_watershed(roi, params, manual_mask)
+                    elif segmentation_method == 'grabcut':
+                        generated_poly, debug_mask = self.get_refined_polygon(roi, params, manual_mask)
+                    elif segmentation_method == 'smart':
+                        print(f"  Calling smart ridge-snapping segmentation...")
+                        generated_poly, debug_mask = self.get_refined_polygon_livewire(roi, params, manual_mask)
+                    else:
+                        generated_poly, debug_mask = self.get_refined_polygon_watershed(roi, params, manual_mask)
+
+                    if len(generated_poly) > 0:
+                        # CRITICAL FIX: Ensure the polygon is stored back in the detection
+                        generated_poly = generated_poly.astype(np.float32)
+                        det['poly'] = generated_poly
+                        det['polygons'] = [generated_poly]
+                        polygons = [generated_poly]
+                        print(f"  Stored smart polygon with {len(generated_poly)} points in detection")
+                        
+                        b, g, r = color
+                        colored_mask = cv2.merge([
+                            np.where(debug_mask > 0, b, 0).astype(np.uint8),
+                            np.where(debug_mask > 0, g, 0).astype(np.uint8),
+                            np.where(debug_mask > 0, r, 0).astype(np.uint8)
+                        ])
+                        draw_img[y1:y2, x1:x2] = cv2.addWeighted(draw_img[y1:y2, x1:x2], 0.6, colored_mask, 0.4, 0)
+
+            # 3. Draw the Polygon and interactive points
+            if show_segments and polygons:
+                for poly_idx, poly_item in enumerate(polygons):
+                    pts = (poly_item.reshape((-1, 1, 2)) + [x1, y1]).astype(np.int32)
+                    cv2.polylines(draw_img, [pts], isClosed=True, color=color, thickness=thickness)
+
+                    if is_selected:
+                        vertex_fill_radius = max(2, base_thickness + 1)
+                        vertex_border_radius = vertex_fill_radius + 1
+                        active_idx = max(0, min(int(selected_island_idx), len(polygons) - 1))
+                        if selected_island_kind == "outer" and poly_idx == active_idx:
+                            for p_idx, p in enumerate(pts):
+                                cv2.circle(draw_img, tuple(p[0]), vertex_fill_radius, (255, 255, 255), -1)
+                                cv2.circle(draw_img, tuple(p[0]), vertex_border_radius, color, max(1, base_thickness - 1))
+                            cv2.putText(
+                                draw_img,
+                                f"outer {active_idx + 1}/{len(polygons)}",
+                                (x1, y2 + 16),
+                                cv2.FONT_HERSHEY_SIMPLEX,
+                                0.45,
+                                color,
+                                1
+                            )
+
+            if show_segments and holes:
+                hole_color = (0, 165, 255)
+                for hole_idx, hole_item in enumerate(holes):
+                    pts = (hole_item.reshape((-1, 1, 2)) + [x1, y1]).astype(np.int32)
+                    cv2.polylines(draw_img, [pts], isClosed=True, color=hole_color, thickness=max(1, thickness - 1))
+                    if is_selected:
+                        vertex_fill_radius = max(2, base_thickness + 1)
+                        vertex_border_radius = vertex_fill_radius + 1
+                        active_idx = max(0, min(int(selected_island_idx), len(holes) - 1))
+                        if selected_island_kind == "hole" and hole_idx == active_idx:
+                            for p in pts:
+                                cv2.circle(draw_img, tuple(p[0]), vertex_fill_radius, (255, 255, 255), -1)
+                                cv2.circle(draw_img, tuple(p[0]), vertex_border_radius, hole_color, max(1, base_thickness - 1))
+                            cv2.putText(
+                                draw_img,
+                                f"hole {active_idx + 1}/{len(holes)}",
+                                (x1, y2 + 32),
+                                cv2.FONT_HERSHEY_SIMPLEX,
+                                0.45,
+                                hole_color,
+                                1
+                            )
+
+            # 4. Visualize Brush Strokes with colored overlay
+            if show_segments and is_selected and manual_mask is not None:
+                # Get the actual ROI region from the image
+                roi_region = draw_img[y1:y2, x1:x2]
+                roi_h, roi_w = roi_region.shape[:2]
+
+                # Ensure manual_mask matches ROI size
+                if manual_mask.shape[:2] != (roi_h, roi_w):
+                    manual_mask_resized = cv2.resize(manual_mask, (roi_w, roi_h), interpolation=cv2.INTER_NEAREST)
+                else:
+                    manual_mask_resized = manual_mask
+
+                # Create colored overlay for brush strokes
+                overlay = np.zeros((roi_h, roi_w, 3), dtype=np.uint8)
+
+                # Background strokes in RED (semi-transparent)
+                bg_mask = manual_mask_resized == 0
+                overlay[bg_mask] = [0, 0, 255]  # Red in BGR
+
+                # Foreground strokes in GREEN (semi-transparent)
+                fg_mask = manual_mask_resized == 1
+                overlay[fg_mask] = [0, 255, 0]  # Green in BGR
+
+                # Only blend where we have brush strokes (not the "probable" areas)
+                mask_combined = (bg_mask | fg_mask).astype(np.uint8)
+                if mask_combined.sum() > 0:
+                    # Create alpha channel effect
+                    alpha = 0.4  # 40% opacity for the overlay
+                    blended = cv2.addWeighted(roi_region, 1 - alpha, overlay, alpha, 0)
+                    # Only apply blending where strokes exist
+                    roi_region[mask_combined > 0] = blended[mask_combined > 0]
+                    draw_img[y1:y2, x1:x2] = roi_region
+
+        return draw_img
+
+    @staticmethod
+    def measure_focus(raw_frame):
+        return cv2.Laplacian(raw_frame, cv2.CV_64F).var()
+
+    @staticmethod
+    def draw_fps_overlay(image, fps_in, fps_out, fps_pred=0.0):
+        """Draw input/output stream FPS overlay on an annotated frame."""
+        if image is None:
+            return image
+        text = f"FPS in (pserv): {fps_in:.1f} | out (unified): {fps_out:.1f} | pred: {fps_pred:.1f}"
+        cv2.putText(
+            image,
+            text,
+            (10, 88),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.6,
+            (0, 255, 255),
+            2
+        )
+        return image
+
+    @staticmethod
+    def get_refined_polygon_livewire(roi, p, manual_mask=None):
+        """
+        Smart selection using Gradient ridge-snapping.
+        """
+        h, w = roi.shape[:2]
+        if h < 10 or w < 10: return np.array([]), np.zeros((h, w), dtype=np.uint8)
+
+        # 1. Edge Detection
+        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+        
+        # High sensitivity Canny
+        low_t = max(10, p.get("high", 150) // 3)
+        edges = cv2.Canny(blurred, low_t, p.get("high", 150))
+        
+        # 2. Use Manual Mask to block areas
+        if manual_mask is not None:
+            if manual_mask.shape[:2] != (h, w):
+                manual_mask = cv2.resize(manual_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+            # Remove edges in background-painted areas
+            edges[manual_mask == 0] = 0
+
+        # 3. Create a 'snapping' surface using distance transform
+        # This creates a 'valley' at every edge
+        inv_edges = cv2.bitwise_not(edges)
+        dist_trans = cv2.distanceTransform(inv_edges, cv2.DIST_L2, 5)
+        
+        # 4. Generate a base mask (Watershed is actually a great seed for this)
+        _, seed_mask = ImageProcessor.get_refined_polygon_watershed(roi, p, manual_mask)
+        
+        # 5. Shrink/Expand seed to snap to the distance transform ridge
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        snapped_mask = cv2.morphologyEx(seed_mask, cv2.MORPH_CLOSE, kernel)
+
+        # 6. Find final contours on the snapped mask
+        # Use CHAIN_APPROX_SIMPLE for reliable point extraction
+        contours, _ = cv2.findContours(snapped_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if not contours:
+            print("  Smart: No contours found in snapped_mask")
+            return np.array([]), snapped_mask
+
+        cnt = max(contours, key=cv2.contourArea)
+        
+        # 7. Apply smoothing and ensure float32 format for storage
+        eps_factor = max(0.0001, p.get("eps", 0.005))
+        epsilon = eps_factor * cv2.arcLength(cnt, True)
+        approx = cv2.approxPolyDP(cnt, epsilon, True)
+
+        # Ensure the polygon is a float32 array in (N, 1, 2) shape
+        poly_out = approx.astype(np.float32)
+        print(f"  Smart: Generated polygon with {len(poly_out)} points")
+
+        return poly_out, snapped_mask

src/core/inference_client.py

@@ -0,0 +1,162 @@
+"""Client for ML inference server communication."""
+import requests
+import cv2
+import json
+import numpy as np
+
+
+class InferenceClient:
+    """Handles communication with ML inference server."""
+
+    def __init__(self, server_url):
+        """
+        Initialize inference client.
+
+        Args:
+            server_url: Base URL of inference server (e.g., "http://mx-ml.psi.ch:8002")
+        """
+        self.server_url = server_url
+
+    def predict(self, image, timeout=10):
+        """
+        Send image to server for prediction.
+
+        Args:
+            image: OpenCV image (BGR format)
+            timeout: Request timeout in seconds
+
+        Returns:
+            List of detection results, or None on error
+
+        Raises:
+            requests.RequestException: On connection errors
+            ValueError: On invalid response
+        """
+        _, buf = cv2.imencode('.jpg', image)
+        url = f"{self.server_url}/predict/"
+        files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
+
+        response = requests.post(url, files=files, timeout=timeout)
+
+        if response.status_code == 200:
+            data = response.json()
+            return data.get('results', [])
+        else:
+            raise ValueError(f"Server returned status {response.status_code}")
+
+    def segment_anything(self, image, x, y, timeout=15):
+        """
+        Request a high-precision mask from the SAM backend.
+        """
+        _, buf = cv2.imencode('.jpg', image)
+        # Use trailing slash to match the server's /sam/ route exactly
+        url = f"{self.server_url.rstrip('/')}/sam/"
+        params = {"point_x": int(x), "point_y": int(y)}
+        files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
+
+        response = requests.post(url, files=files, params=params, timeout=timeout)
+        
+        if response.status_code == 200:
+            return response.json()
+        else:
+            raise ValueError(f"SAM error: {response.status_code} - {response.text}")
+
+    def segment_anything_box(self, image, x1, y1, x2, y2, timeout=15):
+        """
+        Request a high-precision mask from the SAM backend using a bounding box prompt.
+        """
+        _, buf = cv2.imencode('.jpg', image)
+        url = f"{self.server_url.rstrip('/')}/sam_box/"
+        params = {
+            "x1": int(x1),
+            "y1": int(y1),
+            "x2": int(x2),
+            "y2": int(y2)
+        }
+        files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
+
+        response = requests.post(url, files=files, params=params, timeout=timeout)
+
+        if response.status_code == 200:
+            return response.json()
+        else:
+            raise ValueError(f"SAM Box error: {response.status_code} - {response.text}")
+
+    def segment_multi_box(self, image, boxes, timeout=30):
+        """
+        Request masks for multiple boxes in one call.
+        boxes: List of [x1, y1, x2, y2]
+        """
+        _, buf = cv2.imencode('.jpg', image)
+        url = f"{self.server_url.rstrip('/')}/sam_multi_box/"
+        params = {"boxes_json": json.dumps(boxes)}
+        files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
+
+        response = requests.post(url, files=files, params=params, timeout=timeout)
+        if response.status_code == 200:
+            return response.json().get('results', [])
+        else:
+            raise ValueError(f"SAM Multi Box error: {response.status_code}")
+
+    def predict_depth(self, image, timeout=20, colorize=True, depth_model=None, depth_max_side=0):
+        """
+        Request depth map for an image.
+
+        Returns:
+            OpenCV image array decoded from PNG response.
+        """
+        _, buf = cv2.imencode('.jpg', image)
+        url = f"{self.server_url.rstrip('/')}/depth/"
+        params = {"colorize": "true" if colorize else "false"}
+        if depth_model:
+            params["depth_model"] = str(depth_model)
+        if isinstance(depth_max_side, int) and depth_max_side > 0:
+            params["depth_max_side"] = str(depth_max_side)
+        files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
+
+        response = requests.post(url, files=files, params=params, timeout=timeout)
+        if response.status_code != 200:
+            raise ValueError(f"Depth error: {response.status_code} - {response.text}")
+        if "image/png" not in response.headers.get("content-type", ""):
+            try:
+                err = response.json().get("error", response.text)
+            except Exception:
+                err = response.text
+            raise ValueError(f"Depth error: {err}")
+
+        depth_png = cv2.imdecode(np.frombuffer(response.content, np.uint8), cv2.IMREAD_UNCHANGED)
+        if depth_png is None:
+            raise ValueError("Depth response decode failed.")
+        return depth_png
+
+    def send_training_data(self, image, metadata, timeout=10):
+        """
+        Send annotated image to training pipeline.
+
+        Args:
+            image: OpenCV image (BGR format)
+            metadata: Dict containing params and detections
+            timeout: Request timeout in seconds
+
+        Returns:
+            Response object
+
+        Raises:
+            requests.RequestException: On connection errors
+        """
+        _, buf = cv2.imencode('.jpg', image)
+        url = f"{self.server_url}/train/"
+        files = {'file': ('img.jpg', buf.tobytes(), 'image/jpeg')}
+        data = {"metadata": json.dumps(metadata)}
+
+        response = requests.post(url, files=files, data=data, timeout=timeout)
+        response.raise_for_status()
+        return response
+
+    def set_server_url(self, url):
+        """Update server URL."""
+        self.server_url = url
+
+    def get_server_url(self):
+        """Get current server URL."""
+        return self.server_url

src/core/zmq_client.py

@@ -0,0 +1,201 @@
+"""ZMQ client for receiving live image streams with detection results."""
+import zmq
+import json
+import numpy as np
+import cv2
+import threading
+
+
+class ZMQStreamClient:
+    """Handles ZMQ subscription for live detector stream."""
+
+    def __init__(self, zmq_url, on_frame_callback, on_focus_callback=None, use_cuda=True):
+        """
+        Initialize ZMQ stream client.
+
+        Args:
+            zmq_url: ZMQ server URL (e.g., "tcp://localhost:9091")
+            on_frame_callback: Callback function(frame, detections, metadata) called on new frame
+            on_focus_callback: Optional callback function(focus_score) for focus updates
+            use_cuda: Use OpenCV CUDA path when available on this client
+        """
+        self.zmq_url = zmq_url
+        self.on_frame_callback = on_frame_callback
+        self.on_focus_callback = on_focus_callback
+        self.focus_enabled = on_focus_callback is not None
+        self.use_cuda = bool(use_cuda)
+        self.cuda_enabled = False
+        self._cuda_lap_filter = None
+        self.streaming = False
+        self.worker_thread = None
+
+        if self.use_cuda:
+            try:
+                if hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0:
+                    # Create reusable Laplacian filter for focus score on GPU.
+                    self._cuda_lap_filter = cv2.cuda.createLaplacianFilter(cv2.CV_8UC1, cv2.CV_32F, 3)
+                    self.cuda_enabled = True
+            except Exception:
+                self.cuda_enabled = False
+                self._cuda_lap_filter = None
+        self._cuda_fallback_logged = False
+
+    def start(self):
+        """Start streaming in background thread."""
+        if self.streaming:
+            return
+        self.streaming = True
+        self.worker_thread = threading.Thread(target=self._worker, daemon=True)
+        self.worker_thread.start()
+
+    def stop(self):
+        """Stop streaming."""
+        self.streaming = False
+        if self.worker_thread:
+            self.worker_thread.join(timeout=1.0)
+
+    def is_streaming(self):
+        """Check if currently streaming."""
+        return self.streaming
+
+    def enable_focus(self):
+        """Enable focus-score computation and callbacks."""
+        self.focus_enabled = True
+
+    def disable_focus(self):
+        """Disable focus-score computation and callbacks."""
+        self.focus_enabled = False
+
+    def _worker(self):
+        """Background worker that listens to ZMQ stream."""
+        print(f"DEBUG: Starting Unified ZMQ worker ({self.zmq_url}) cuda={self.cuda_enabled}", flush=True)
+        if self.use_cuda:
+            if self.cuda_enabled:
+                print("INFO: OpenCV CUDA path enabled for focus scoring", flush=True)
+            else:
+                print("INFO: OpenCV CUDA requested but unavailable, using CPU path", flush=True)
+        ctx = zmq.Context()
+        sub = ctx.socket(zmq.SUB)
+        sub.connect(self.zmq_url)
+        sub.setsockopt_string(zmq.SUBSCRIBE, "")
+        sub.setsockopt(zmq.RCVBUF, 128 * 1024 * 1024)
+        poller = zmq.Poller()
+        poller.register(sub, zmq.POLLIN)
+
+        while self.streaming:
+            try:
+                socks = dict(poller.poll(10))
+                if sub in socks:
+                    # Drain queue to get latest message
+                    latest_msg = None
+                    while True:
+                        try:
+                            latest_msg = sub.recv_multipart(zmq.NOBLOCK)
+                        except zmq.Again:
+                            break
+
+                    if latest_msg:
+                        self._process_message(latest_msg)
+            except Exception as e:
+                print(f"ZMQ Error: {e}")
+                break
+
+        sub.close()
+        ctx.term()
+
+    def _process_message(self, msg_parts):
+        """Process received ZMQ multipart message."""
+        try:
+            # Last part is image data, rest might contain header
+            img_data = msg_parts[-1]
+            header = None
+
+            # Try to find JSON header in message parts
+            for part in msg_parts[:-1]:
+                try:
+                    if not part:
+                        continue
+                    header = json.loads(part.decode('utf-8'))
+                    if isinstance(header, dict) and 'shape' in header:
+                        break
+                    header = None
+                except:
+                    continue
+
+            if not header:
+                return
+
+            # Extract detections
+            detections = header.get('boxes', [])
+            if isinstance(detections, list):
+                for det in detections:
+                    if not isinstance(det, dict):
+                        continue
+                    poly = det.get('poly')
+                    if poly is None:
+                        continue
+                    try:
+                        poly_arr = np.asarray(poly, dtype=np.float32)
+                        if poly_arr.ndim == 2 and poly_arr.shape[1] == 2 and poly_arr.shape[0] >= 3:
+                            det['poly'] = poly_arr.reshape(-1, 1, 2)
+                        else:
+                            det.pop('poly', None)
+                    except Exception:
+                        det.pop('poly', None)
+
+            # Decode image payload (raw BGR or Bayer from unified stream server).
+            pixel_format = str(header.get('pixel_format', 'BGR')).upper()
+            shape = header.get('shape', [])
+            if pixel_format == 'BAYER':
+                if len(shape) != 2:
+                    return
+                h, w = int(shape[0]), int(shape[1])
+                raw = np.frombuffer(img_data, np.uint8).reshape((h, w))
+                bayer_pattern = str(header.get('bayer_pattern', 'RGGB')).upper()
+                demosaic_map = {
+                    'RGGB': cv2.COLOR_BAYER_RG2BGR,
+                    'GBRG': cv2.COLOR_BAYER_GB2BGR,
+                    'GRBG': cv2.COLOR_BAYER_GR2BGR,
+                    'BGGR': cv2.COLOR_BAYER_BG2BGR,
+                }
+                demosaic_code = demosaic_map.get(bayer_pattern)
+                if demosaic_code is None:
+                    return
+                img = cv2.cvtColor(raw, demosaic_code)
+            else:
+                if len(shape) != 3:
+                    return
+                h, w, c = int(shape[0]), int(shape[1]), int(shape[2])
+                if c != 3:
+                    return
+                img = np.frombuffer(img_data, np.uint8).reshape((h, w, c))
+
+            # Calculate focus score
+            if self.focus_enabled and self.on_focus_callback:
+                if self.cuda_enabled:
+                    try:
+                        gpu_img = cv2.cuda_GpuMat()
+                        gpu_img.upload(img)
+                        gpu_gray = cv2.cuda.cvtColor(gpu_img, cv2.COLOR_BGR2GRAY)
+                        gpu_lap = self._cuda_lap_filter.apply(gpu_gray)
+                        lap = gpu_lap.download()
+                        focus_score = float(lap.var())
+                    except Exception:
+                        if not self._cuda_fallback_logged:
+                            print("WARN: CUDA focus path failed at runtime, falling back to CPU", flush=True)
+                            self._cuda_fallback_logged = True
+                        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+                        focus_score = cv2.Laplacian(gray, cv2.CV_64F).var()
+                else:
+                    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+                    focus_score = cv2.Laplacian(gray, cv2.CV_64F).var()
+                self.on_focus_callback(focus_score)
+
+            # Call frame callback. Keep backward compatibility with 2-arg callbacks.
+            try:
+                self.on_frame_callback(img, detections, header)
+            except TypeError:
+                self.on_frame_callback(img, detections)
+
+        except Exception as e:
+            print(f"Error processing ZMQ message: {e}")

src/gui/__init__.py

@@ -0,0 +1,5 @@
+"""GUI components for AareLC ML Studio."""
+from .canvas_panel import MainCanvas
+from .control_panel import ControlPanel
+
+__all__ = ['MainCanvas', 'ControlPanel']

src/gui/canvas_interaction_mixin.py

@@ -0,0 +1,697 @@
+"""Canvas drawing and interaction mixin for the GUI."""
+
+import math
+
+import cv2
+import numpy as np
+
+
+class CanvasInteractionMixin:
+    """Canvas rendering and user interaction logic."""
+
+    @staticmethod
+    def _clamp_point(pt, width, height):
+        if pt is None:
+            return None
+        x = max(0, min(width - 1, int(pt[0])))
+        y = max(0, min(height - 1, int(pt[1])))
+        return x, y
+
+    def _extract_target_from_metadata(self, metadata):
+        if not isinstance(metadata, dict):
+            return None, None
+        point = metadata.get('target_point')
+        if not isinstance(point, dict):
+            return None, None
+        x = point.get('x')
+        y = point.get('y')
+        if not isinstance(x, (int, float)) or not isinstance(y, (int, float)):
+            return None, None
+        label = str(point.get('source', 'target_point'))
+        return (int(round(x)), int(round(y))), label
+
+    @staticmethod
+    def _draw_target_cross(image, point, label=None):
+        if image is None or point is None:
+            return image
+        h, w = image.shape[:2]
+        cx = max(0, min(w - 1, int(point[0])))
+        cy = max(0, min(h - 1, int(point[1])))
+        arm = 14
+        color = (255, 0, 255)  # magenta
+        cv2.line(image, (cx - arm, cy), (cx + arm, cy), color, 2)
+        cv2.line(image, (cx, cy - arm), (cx, cy + arm), color, 2)
+        cv2.circle(image, (cx, cy), 3, color, -1)
+        if label:
+            cv2.putText(
+                image,
+                f"Target: {label} ({cx},{cy})",
+                (max(8, cx + 10), max(20, cy - 10)),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                color,
+                1
+            )
+        return image
+
+    def _toggle_analysis_grid(self, event=None):
+        self.analysis_state['show_grid'] = not self.analysis_state['show_grid']
+        self.redraw_current_frame()
+        self.status_var.set(f"Measure grid: {'ON' if self.analysis_state['show_grid'] else 'OFF'}")
+
+    def _toggle_analysis_crosshair(self, event=None):
+        self.analysis_state['show_crosshair'] = not self.analysis_state['show_crosshair']
+        self.redraw_current_frame()
+        self.status_var.set(f"Measure crosshair: {'ON' if self.analysis_state['show_crosshair'] else 'OFF'}")
+
+    def _clear_measurements(self, event=None):
+        self.analysis_state['line_start'] = None
+        self.analysis_state['line_end'] = None
+        self.analysis_state['line_dragging'] = False
+        self.analysis_state['roi_start'] = None
+        self.analysis_state['roi_end'] = None
+        self.analysis_state['roi_dragging'] = False
+        self.redraw_current_frame()
+        self.status_var.set("Measurement overlays cleared")
+
+    def _draw_analysis_overlay(self, annotated, source_frame):
+        h, w = annotated.shape[:2]
+        state = self.analysis_state
+
+        if state['show_grid']:
+            step = max(40, min(w, h) // 12)
+            for gx in range(step, w, step):
+                cv2.line(annotated, (gx, 0), (gx, h - 1), (70, 70, 70), 1)
+            for gy in range(step, h, step):
+                cv2.line(annotated, (0, gy), (w - 1, gy), (70, 70, 70), 1)
+
+        cursor = self._clamp_point(state['cursor'], w, h)
+        if cursor is not None and state['show_crosshair']:
+            cx, cy = cursor
+            cv2.line(annotated, (cx, 0), (cx, h - 1), (0, 180, 255), 1)
+            cv2.line(annotated, (0, cy), (w - 1, cy), (0, 180, 255), 1)
+
+        if cursor is not None:
+            cx, cy = cursor
+            b, g, r = source_frame[cy, cx]
+            cv2.putText(
+                annotated,
+                f"Cursor ({cx},{cy}) BGR=({int(b)},{int(g)},{int(r)})",
+                (10, h - 14),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                (50, 220, 255),
+                1
+            )
+
+        line_start = self._clamp_point(state['line_start'], w, h)
+        line_end = self._clamp_point(state['line_end'], w, h)
+        if line_start is not None and line_end is not None:
+            cv2.line(annotated, line_start, line_end, (0, 255, 255), 2)
+            cv2.circle(annotated, line_start, 4, (0, 255, 255), -1)
+            cv2.circle(annotated, line_end, 4, (0, 255, 255), -1)
+            dx = line_end[0] - line_start[0]
+            dy = line_end[1] - line_start[1]
+            dist = math.hypot(dx, dy)
+            angle_deg = math.degrees(math.atan2(dy, dx))
+            cv2.putText(
+                annotated,
+                f"Line: {dist:.1f}px  dx={dx} dy={dy}  angle={angle_deg:.1f}deg",
+                (10, 44),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.55,
+                (0, 255, 255),
+                2
+            )
+
+        roi_start = self._clamp_point(state['roi_start'], w, h)
+        roi_end = self._clamp_point(state['roi_end'], w, h)
+        if roi_start is not None and roi_end is not None:
+            x1, x2 = sorted((roi_start[0], roi_end[0]))
+            y1, y2 = sorted((roi_start[1], roi_end[1]))
+            cv2.rectangle(annotated, (x1, y1), (x2, y2), (255, 255, 0), 2)
+            roi_w = x2 - x1
+            roi_h = y2 - y1
+            roi_area = roi_w * roi_h
+
+            roi_label = f"ROI: {roi_w}x{roi_h}px area={roi_area}px2"
+            if roi_w > 1 and roi_h > 1:
+                roi = source_frame[y1:y2, x1:x2]
+                if roi.size > 0:
+                    mean_bgr = roi.reshape(-1, 3).mean(axis=0)
+                    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
+                    roi_label += (
+                        f" gray={float(gray.mean()):.1f}+/-{float(gray.std()):.1f}"
+                        f" meanBGR=({mean_bgr[0]:.1f},{mean_bgr[1]:.1f},{mean_bgr[2]:.1f})"
+                    )
+
+            cv2.putText(
+                annotated,
+                roi_label,
+                (10, 66),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                (255, 255, 0),
+                1
+            )
+
+        if self.tk_vars['active_tool'].get() == "measure":
+            cv2.putText(
+                annotated,
+                "Measure tool: L-drag line | R-drag ROI | g grid | c crosshair | x clear",
+                (10, h - 34),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                (220, 220, 220),
+                1
+            )
+        return annotated
+
+    def redraw_current_frame(self):
+        """Redraw current frame with annotations."""
+        if self.current_raw_frame is not None:
+            self.update_canvas(self.current_raw_frame)
+
+    def update_canvas(self, frame, metadata=None):
+        """Update canvas with annotated frame."""
+        if frame is None:
+            return
+
+        is_streaming = self.zmq_client is not None and self.zmq_client.is_streaming()
+
+        # Pass the selected index to the processor for highlighting
+        selected_idx = getattr(self.canvas, 'selected_obj_idx', -1)
+
+        # Get hidden classes from control panel
+        hidden_classes = {
+            cls_name for cls_name, var in self.ctrl_panel.visibility_vars.items()
+            if not var.get()
+        }
+
+        # Compute edge map if not already done or image changed
+        if self.image_processor.edge_map is None or self.image_processor.edge_map_shape != frame.shape[:2]:
+            self.image_processor.compute_edge_map(frame)
+
+        annotated = self.image_processor.draw_detections(
+            frame,
+            self.last_results,
+            self.class_params,
+            self.class_colors,
+            active_class=self.tk_vars['active_class'].get(),
+            show_debug_edges=self.tk_vars['show_debug_edges'].get(),
+            refine_mode=not is_streaming,
+            selected_idx=selected_idx,
+            selected_island_idx=self.selected_island_idx,
+            selected_island_kind=self.selected_island_kind,
+            hidden_classes=hidden_classes,
+            segmentation_method=self.tk_vars['segmentation_method'].get(),
+            annotation_line_width=self.tk_vars['annotation_line_width'].get(),
+            show_boxes=self.tk_vars['show_boxes'].get(),
+            show_segments=self.tk_vars['show_segments'].get(),
+        )
+        if self.tk_vars['active_tool'].get() == "measure":
+            annotated = self._draw_analysis_overlay(annotated, frame)
+        if is_streaming:
+            annotated = self.image_processor.draw_fps_overlay(
+                annotated,
+                self.stream_fps_in,
+                self.stream_fps_out,
+                self.stream_fps_pred
+            )
+        point, label = self._extract_target_from_metadata(metadata if metadata is not None else self._last_stream_metadata)
+        annotated = self._draw_target_cross(annotated, point, label)
+        self.canvas.display_image(annotated)
+
+    def _get_hidden_classes(self):
+        """Get set of currently hidden classes."""
+        return {
+            cls_name for cls_name, var in self.ctrl_panel.visibility_vars.items()
+            if not var.get()
+        }
+
+    def _on_canvas_interaction(self, click_event=None, drag_event=None, right_click_event=None,
+                               double_click_event=None, move_event=None, right_drag_event=None,
+                               right_release_event=None):
+        if self.current_raw_frame is None:
+            return
+        img_h, img_w = self.current_raw_frame.shape[:2]
+        active_tool = self.tk_vars['active_tool'].get()
+        self.canvas.suspend_pan = (active_tool == "measure")
+
+        if active_tool != "measure" and (move_event is not None or right_drag_event is not None or right_release_event is not None):
+            return
+
+        if active_tool == "measure":
+            state = self.analysis_state
+
+            # End left-drag measurement when canvas signals release (no event payload)
+            if not any([click_event, drag_event, right_click_event, double_click_event, move_event, right_drag_event, right_release_event]):
+                state['line_dragging'] = False
+                return
+
+            active_event = move_event or click_event or right_click_event or double_click_event or right_drag_event
+            if drag_event:
+                active_event = drag_event[0]
+            if right_release_event:
+                active_event = right_release_event
+
+            if active_event is not None:
+                ix, iy = self.canvas.get_image_coords(active_event.x, active_event.y, img_w, img_h)
+                state['cursor'] = (ix, iy)
+
+            if click_event is not None:
+                ix, iy = self.canvas.get_image_coords(click_event.x, click_event.y, img_w, img_h)
+                state['line_start'] = (ix, iy)
+                state['line_end'] = (ix, iy)
+                state['line_dragging'] = True
+
+            if drag_event is not None and state['line_dragging']:
+                event, _, _ = drag_event
+                ix, iy = self.canvas.get_image_coords(event.x, event.y, img_w, img_h)
+                state['line_end'] = (ix, iy)
+
+            if right_click_event is not None:
+                ix, iy = self.canvas.get_image_coords(right_click_event.x, right_click_event.y, img_w, img_h)
+                state['roi_start'] = (ix, iy)
+                state['roi_end'] = (ix, iy)
+                state['roi_dragging'] = True
+
+            if right_drag_event is not None and state['roi_dragging']:
+                ix, iy = self.canvas.get_image_coords(right_drag_event.x, right_drag_event.y, img_w, img_h)
+                state['roi_end'] = (ix, iy)
+
+            if right_release_event is not None:
+                ix, iy = self.canvas.get_image_coords(right_release_event.x, right_release_event.y, img_w, img_h)
+                state['roi_end'] = (ix, iy)
+                state['roi_dragging'] = False
+
+            self.redraw_current_frame()
+            return
+
+        # --- SAM CLICK LOGIC ---
+        if click_event and active_tool == "sam":
+            ix, iy = self.canvas.get_image_coords(click_event.x, click_event.y, img_w, img_h)
+            self._run_sam_inference(ix, iy)
+            return
+
+        # Initialize variables
+        hidden_classes = self._get_hidden_classes()
+        brush_mode = self.tk_vars['brush_mode'].get()
+        event_obj = click_event or right_click_event or double_click_event or (drag_event[0] if drag_event else None)
+
+        is_opt = False
+        if event_obj:
+            is_opt = (event_obj.state & 0x0008) or (event_obj.state & 0x0010)
+
+        # --- BRUSH LOGIC (Handles both Click and Drag) ---
+        # Check BOTH old active_tool and new brush_mode checkbox
+        is_brush_active = (active_tool == "brush" or brush_mode or (right_click_event and is_opt))
+
+        if event_obj and is_brush_active:
+            ix, iy = self.canvas.get_image_coords(event_obj.x, event_obj.y, img_w, img_h)
+
+            if self.canvas.selected_obj_idx != -1:
+                det = self.last_results[self.canvas.selected_obj_idx]
+                box = det.get('box', det)
+                x1, y1 = int(box['x1']), int(box['y1'])
+                x2, y2 = int(box['x2']), int(box['y2'])
+
+                # Save history on first brush stroke
+                if not getattr(self, '_brush_started', False):
+                    self._save_history()
+                    self._brush_started = True
+
+                # Get brush parameters
+                brush_size = self.tk_vars['brush_size'].get()
+                brush_type = self.tk_vars['brush_type'].get()
+
+                # Convert to local coordinates
+                lx, ly = int(ix - x1), int(iy - y1)
+
+                if 0 <= lx < (x2 - x1) and 0 <= ly < (y2 - y1):
+                    # Create or get the manual mask
+                    if 'manual_mask' not in det:
+                        h, w = y2 - y1, x2 - x1
+                        det['manual_mask'] = np.full((h, w), 2, dtype=np.uint8)  # 2 = probable foreground
+
+                    # Paint with circular brush
+                    mask_value = 0 if brush_type == 'background' else 1  # 0=BGD, 1=FGD
+
+                    # Draw circular brush stroke on the mask
+                    cv2.circle(det['manual_mask'], (lx, ly), brush_size, mask_value, -1)
+
+                    # Don't auto-delete polygon - wait for manual regeneration
+                    # This prevents slowdown from constant regeneration during painting
+
+                    self.redraw_current_frame()
+                    self.status_var.set(f"Brush: {'Background' if brush_type == 'background' else 'Foreground'} region marked - Click 'Regenerate Edge' to update")
+                    return  # Stop other interactions while brushing
+
+        # --- REGULAR INTERACTION (Select Tool) ---
+        active_event = click_event or right_click_event or double_click_event
+        if active_event:
+            ix, iy = self.canvas.get_image_coords(active_event.x, active_event.y, img_w, img_h)
+
+            # --- DOUBLE CLICK: Add/Insert Point ---
+            if double_click_event and self.canvas.selected_obj_idx != -1:
+                self._save_history()
+                det = self.last_results[self.canvas.selected_obj_idx]
+                box = det.get('box', det)
+                poly = self._get_active_polygon(det)
+                lx, ly = ix - box.get('x1', 0), iy - box.get('y1', 0)
+
+                if poly is None or len(poly) < 2:
+                    new_poly = np.array([[[lx, ly]], [[lx + 50, ly]], [[lx + 50, ly + 50]]], dtype=np.float32)
+                    if self.selected_island_kind == "hole":
+                        holes = self._det_get_holes_relative(det)
+                        if holes:
+                            holes[self.selected_island_idx] = new_poly
+                            self._set_det_holes(det, holes)
+                        else:
+                            self._set_det_holes(det, [new_poly])
+                    else:
+                        polygons = self._det_get_polygons_relative(det)
+                        if polygons:
+                            polygons[self.selected_island_idx] = new_poly
+                            self._set_det_polygons(det, polygons)
+                        else:
+                            self._set_det_single_polygon(det, new_poly)
+                else:
+                    best_dist, insert_idx = float('inf'), len(poly)
+                    for i in range(len(poly)):
+                        p1, p2 = poly[i].reshape(-1), poly[(i + 1) % len(poly)].reshape(-1)
+                        d = np.linalg.norm(p1 - [lx, ly]) + np.linalg.norm(p2 - [lx, ly])
+                        if d < best_dist:
+                            best_dist, insert_idx = d, i + 1
+                    updated_poly = np.insert(poly, insert_idx, [[[lx, ly]]], axis=0)
+                    if self.selected_island_kind == "hole":
+                        holes = self._det_get_holes_relative(det)
+                        if holes:
+                            holes[self.selected_island_idx] = updated_poly
+                            self._set_det_holes(det, holes)
+                        else:
+                            self._set_det_holes(det, [updated_poly])
+                    else:
+                        polygons = self._det_get_polygons_relative(det)
+                        if polygons:
+                            polygons[self.selected_island_idx] = updated_poly
+                            self._set_det_polygons(det, polygons)
+                        else:
+                            self._set_det_single_polygon(det, updated_poly)
+                return self.redraw_current_frame()
+
+            # --- HIT TESTING: Select or Delete ---
+            hit_radius = 15 / self.canvas.zoom_level
+
+            # Hit Test: Box corners FIRST so resize handles win over nearby polygon points
+            for i, det in enumerate(self.last_results):
+                if det.get('name', det.get('label', 'unknown')) in hidden_classes:
+                    continue
+                box = det.get('box', det)
+                x1, y1 = box.get('x1', 0), box.get('y1', 0)
+                x2, y2 = box.get('x2', 0), box.get('y2', 0)
+                corners = {
+                    'box_resize_tl': (x1, y1),
+                    'box_resize_tr': (x2, y1),
+                    'box_resize_bl': (x1, y2),
+                    'box_resize_br': (x2, y2),
+                }
+                for mode, (cx, cy) in corners.items():
+                    if np.linalg.norm(np.array([cx, cy]) - [ix, iy]) < hit_radius:
+                        self.canvas.selected_obj_idx = i
+                        self.selected_island_idx = 0
+                        self.selected_island_kind = "outer"
+                        self.canvas.drag_mode = mode
+                        self.canvas.selected_point_idx = -1
+                        return self.redraw_current_frame()
+
+            # Hit Test: Points (after corners)
+            for i, det in enumerate(self.last_results):
+                if det.get('name', det.get('label', 'unknown')) in hidden_classes:
+                    continue
+                box = det.get('box', det)
+                rings = [("outer", self._det_get_polygons_relative(det)), ("hole", self._det_get_holes_relative(det))]
+                for ring_kind, ring_list in rings:
+                    for island_idx, poly in enumerate(ring_list):
+                        global_pts = poly.reshape(-1, 2) + [box.get('x1', 0), box.get('y1', 0)]
+                        for p_idx, pt in enumerate(global_pts):
+                            if np.linalg.norm(pt - [ix, iy]) < hit_radius:
+                                self.selected_island_kind = ring_kind
+                                self.selected_island_idx = island_idx
+                                if right_click_event:
+                                    self._save_history()
+                                    updated = np.delete(poly, p_idx, axis=0)
+                                    ring_list[island_idx] = updated
+                                    if ring_kind == "outer":
+                                        self._set_det_polygons(det, ring_list)
+                                    else:
+                                        self._set_det_holes(det, ring_list)
+                                    return self.redraw_current_frame()
+                                self.canvas.selected_obj_idx, self.canvas.selected_point_idx = i, p_idx
+                                self.canvas.drag_mode = 'point_move'
+                                return self.redraw_current_frame()
+
+            # Hit Test: polygon interior to select active island
+            for i, det in enumerate(self.last_results):
+                if det.get('name', det.get('label', 'unknown')) in hidden_classes:
+                    continue
+                box = det.get('box', det)
+                rings = [("hole", self._det_get_holes_relative(det)), ("outer", self._det_get_polygons_relative(det))]
+                for ring_kind, ring_list in rings:
+                    for island_idx, poly in enumerate(ring_list):
+                        global_pts = (poly.reshape(-1, 2) + [box.get('x1', 0), box.get('y1', 0)]).astype(np.float32)
+                        inside = cv2.pointPolygonTest(global_pts.reshape(-1, 1, 2), (float(ix), float(iy)), False)
+                        if inside >= 0:
+                            self.canvas.selected_obj_idx = i
+                            self.selected_island_kind = ring_kind
+                            self.selected_island_idx = island_idx
+                            self.canvas.selected_point_idx = -1
+                            self.canvas.drag_mode = 'box_move'
+                            return self.redraw_current_frame()
+
+            # Hit Test: Box body with overlap disambiguation (top-most first + click-to-cycle)
+            hit_indices = []
+            for i in range(len(self.last_results) - 1, -1, -1):
+                det = self.last_results[i]
+                if det.get('name', det.get('label', 'unknown')) in hidden_classes:
+                    continue
+                box = det.get('box', det)
+                x1, y1, x2, y2 = box.get('x1', 0), box.get('y1', 0), box.get('x2', 0), box.get('y2', 0)
+                if x1 < ix < x2 and y1 < iy < y2:
+                    hit_indices.append(i)
+
+            if hit_indices:
+                if right_click_event and not is_opt:
+                    target_idx = hit_indices[0]
+                    self._save_history()
+                    self.last_results.pop(target_idx)
+                    self.canvas.selected_obj_idx = -1
+                    self.selected_island_idx = 0
+                    self.selected_island_kind = "outer"
+                    return self.redraw_current_frame()
+
+                selected_idx = hit_indices[0]
+                if click_event and len(hit_indices) > 1:
+                    cycle_key = tuple(hit_indices)
+                    pick = self._last_overlap_pick
+                    same_spot = False
+                    if pick is not None:
+                        px, py = pick.get("point", (None, None))
+                        same_spot = px is not None and py is not None and abs(px - ix) <= 6 and abs(py - iy) <= 6
+                    if pick and pick.get("key") == cycle_key and same_spot:
+                        selected_idx = hit_indices[(pick.get("pos", 0) + 1) % len(hit_indices)]
+                    pos = hit_indices.index(selected_idx)
+                    self._last_overlap_pick = {"key": cycle_key, "pos": pos, "point": (ix, iy)}
+                else:
+                    self._last_overlap_pick = {"key": tuple(hit_indices), "pos": 0, "point": (ix, iy)}
+
+                self.canvas.selected_obj_idx, self.canvas.drag_mode, self.canvas.selected_point_idx = selected_idx, 'box_move', -1
+
+                det = self.last_results[selected_idx]
+                self.selected_island_idx = 0
+                self.selected_island_kind = "outer"
+                det_class = det.get('name', det.get('label', self._default_class_name()))
+                if det_class in self.class_params:
+                    self.tk_vars['active_class'].set(det_class)
+
+                return self.redraw_current_frame()
+
+            if click_event:
+                self.canvas.selected_obj_idx = -1
+                self.selected_island_idx = 0
+                self.selected_island_kind = "outer"
+                self.redraw_current_frame()
+
+        # 3. DRAG HANDLING
+        if drag_event:
+            event, dx, dy = drag_event
+            ix, iy = self.canvas.get_image_coords(event.x, event.y, img_w, img_h)
+
+            # Continuous Brush Painting
+            if (active_tool == "brush" or brush_mode or (event.state & 0x0400 and is_opt)) and self.canvas.selected_obj_idx != -1:
+                det = self.last_results[self.canvas.selected_obj_idx]
+                box = det.get('box', det)
+                x1, y1 = int(box['x1']), int(box['y1'])
+                x2, y2 = int(box['x2']), int(box['y2'])
+
+                # Initialize manual mask if needed
+                if 'manual_mask' not in det:
+                    h, w = y2 - y1, x2 - x1
+                    det['manual_mask'] = np.full((h, w), 2, dtype=np.uint8)
+
+                lx, ly = int(ix - x1), int(iy - y1)
+                if 0 <= lx < det['manual_mask'].shape[1] and 0 <= ly < det['manual_mask'].shape[0]:
+                    # Get brush parameters
+                    brush_size = self.tk_vars['brush_size'].get()
+                    brush_type = self.tk_vars['brush_type'].get()
+                    mask_value = 0 if brush_type == 'background' else 1
+
+                    cv2.circle(det['manual_mask'], (lx, ly), brush_size, mask_value, -1)
+                    self.redraw_current_frame()
+                return
+
+            if not getattr(self, '_is_dragging', False):
+                self._save_history()
+                self._is_dragging = True
+
+            idx = self.canvas.selected_obj_idx
+            if idx == -1:
+                return
+            det = self.last_results[idx]
+            box = det.get('box', det)
+
+            if self.canvas.drag_mode == 'point_move':
+                p_idx = self.canvas.selected_point_idx
+                if self.selected_island_kind == "hole":
+                    rings = self._det_get_holes_relative(det)
+                else:
+                    rings = self._det_get_polygons_relative(det)
+                if not rings:
+                    return
+                island_idx = int(max(0, min(self.selected_island_idx, len(rings) - 1)))
+                poly = rings[island_idx].astype(np.float32)
+                poly_view = poly[p_idx]
+                px, py = (poly_view[0][0], poly_view[0][1]) if poly_view.ndim > 1 else (poly_view[0], poly_view[1])
+                nx, ny = px + dx, py + dy
+                if event.state & 0x0001:  # Shift Snap
+                    nx, ny = self.image_processor.find_nearest_edge(nx + box['x1'], ny + box['y1'])
+                    nx -= box['x1']
+                    ny -= box['y1']
+                if poly_view.ndim > 1:
+                    poly_view[0][0], poly_view[0][1] = nx, ny
+                else:
+                    poly_view[0], poly_view[1] = nx, ny
+                rings[island_idx] = poly
+                if self.selected_island_kind == "hole":
+                    self._set_det_holes(det, rings)
+                else:
+                    self._set_det_polygons(det, rings)
+                self.redraw_current_frame()
+            elif 'box_resize' in self.canvas.drag_mode or self.canvas.drag_mode == 'box_move':
+                old_x1, old_y1 = float(box['x1']), float(box['y1'])
+                old_x2, old_y2 = float(box['x2']), float(box['y2'])
+                old_w = max(1.0, old_x2 - old_x1)
+                old_h = max(1.0, old_y2 - old_y1)
+
+                if self.canvas.drag_mode == 'box_resize_br':
+                    box['x2'] += dx
+                    box['y2'] += dy
+                elif self.canvas.drag_mode == 'box_resize_tl':
+                    box['x1'] += dx
+                    box['y1'] += dy
+                elif self.canvas.drag_mode == 'box_resize_tr':
+                    box['x2'] += dx
+                    box['y1'] += dy
+                elif self.canvas.drag_mode == 'box_resize_bl':
+                    box['x1'] += dx
+                    box['y2'] += dy
+                elif self.canvas.drag_mode == 'box_move':
+                    box['x1'] += dx
+                    box['x2'] += dx
+                    box['y1'] += dy
+                    box['y2'] += dy
+                    self.image_processor.compute_edge_map(self.current_raw_frame, box)
+
+                # Keep corners ordered and inside image bounds.
+                min_size = 2.0
+                x1_new, x2_new = sorted((float(box['x1']), float(box['x2'])))
+                y1_new, y2_new = sorted((float(box['y1']), float(box['y2'])))
+                x1_new = max(0.0, min(x1_new, img_w - min_size))
+                y1_new = max(0.0, min(y1_new, img_h - min_size))
+                x2_new = max(x1_new + min_size, min(x2_new, float(img_w)))
+                y2_new = max(y1_new + min_size, min(y2_new, float(img_h)))
+                box['x1'], box['y1'], box['x2'], box['y2'] = x1_new, y1_new, x2_new, y2_new
+
+                # Resize existing polygon together with the box.
+                if self.canvas.drag_mode.startswith('box_resize'):
+                    new_w = max(1.0, float(box['x2']) - float(box['x1']))
+                    new_h = max(1.0, float(box['y2']) - float(box['y1']))
+                    sx = new_w / old_w
+                    sy = new_h / old_h
+                    polygons_rel = self._det_get_polygons_relative(det)
+                    if polygons_rel:
+                        scaled = []
+                        for poly in polygons_rel:
+                            arr = poly.astype(np.float32)
+                            arr[:, 0, 0] *= sx
+                            arr[:, 0, 1] *= sy
+                            scaled.append(arr)
+                        self._set_det_polygons(det, scaled)
+                    holes_rel = self._det_get_holes_relative(det)
+                    if holes_rel:
+                        scaled_holes = []
+                        for hole in holes_rel:
+                            arr = hole.astype(np.float32)
+                            arr[:, 0, 0] *= sx
+                            arr[:, 0, 1] *= sy
+                            scaled_holes.append(arr)
+                        self._set_det_holes(det, scaled_holes)
+
+                self.redraw_current_frame()
+
+            # Reset flags when drag ends
+        if not drag_event:
+            self._is_dragging = False
+            self._brush_started = False
+
+    def add_new_box(self):
+        """Adds a box at the center of the current screen view."""
+        self._save_history()
+        if self.current_raw_frame is None:
+            return
+
+        img_h, img_w = self.current_raw_frame.shape[:2]
+        cls = self.tk_vars['active_class'].get()
+
+        # Get image coordinates for the center of the canvas
+        cw = self.canvas.winfo_width() / 2
+        ch = self.canvas.winfo_height() / 2
+        ix, iy = self.canvas.get_image_coords(cw, ch, img_w, img_h)
+
+        # Create a 150x150 box centered on the view
+        size = 75
+        x1, y1 = max(0, ix - size), max(0, iy - size)
+        x2, y2 = min(img_w, ix + size), min(img_h, iy + size)
+
+        new_det = {
+            "box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
+            "name": cls,
+            # Keep brand-new boxes as rectangles until user explicitly refines/regenerates.
+            "suppress_auto_polygon": True,
+        }
+
+        self.last_results.append(new_det)
+        self.canvas.selected_obj_idx = len(self.last_results) - 1
+        self.selected_island_idx = 0
+        self.selected_island_kind = "outer"
+        self.redraw_current_frame()
+        self.status_var.set("Added new box")
+
+    def delete_selected(self, event=None):
+        if self.canvas.selected_obj_idx != -1:
+            self._save_history()
+            self.last_results.pop(self.canvas.selected_obj_idx)
+            self.canvas.selected_obj_idx = -1
+            self.selected_island_idx = 0
+            self.selected_island_kind = "outer"
+            self.redraw_current_frame()

src/gui/canvas_panel.py

@@ -0,0 +1,144 @@
+"""Main canvas panel for image display with zoom and pan support."""
+import tkinter as tk
+from PIL import Image, ImageTk
+import cv2
+
+class MainCanvas(tk.Canvas):
+    """Canvas widget for displaying images with zoom and pan capabilities."""
+
+    def __init__(self, parent):
+        # 1. Initialize state variables BEFORE super().__init__
+        self.zoom_level = 1.0
+        self.pan_x = 0
+        self.pan_y = 0
+        self.selected_obj_idx = -1
+        self.selected_point_idx = -1
+        self.drag_mode = None
+        self.last_mouse_x = 0
+        self.last_mouse_y = 0
+        self.display_img = None
+        self.on_redraw_callback = None
+        self.suspend_pan = False
+
+        super().__init__(parent, bg="#111", highlightthickness=0)
+
+        # Bindings
+        self.bind("<ButtonPress-1>", self._on_press)
+        self.bind("<Double-Button-1>", self._on_double_click)
+        self.bind("<B1-Motion>", self._on_drag)
+        self.bind("<ButtonRelease-1>", self._on_release)
+        self.bind("<Motion>", self._on_motion)
+
+        # Map both common right-click buttons
+        self.bind("<Button-2>", self._on_right_click)  # macOS
+        self.bind("<Button-3>", self._on_right_click)  # Windows/Linux
+        self.bind("<B2-Motion>", self._on_right_drag)  # macOS
+        self.bind("<B3-Motion>", self._on_right_drag)  # Windows/Linux
+        self.bind("<ButtonRelease-2>", self._on_right_release)  # macOS
+        self.bind("<ButtonRelease-3>", self._on_right_release)  # Windows/Linux
+
+        self.bind("<MouseWheel>", self._handle_zoom)
+        self.bind("<Button-4>", self._handle_zoom)
+        self.bind("<Button-5>", self._handle_zoom)
+
+    def _on_double_click(self, event):
+        """Handle double-click to add a new point to the selected polygon."""
+        if self.on_redraw_callback:
+            self.on_redraw_callback(double_click_event=event)
+
+    def _on_right_click(self, event):
+        """Handle right-click event for point deletion."""
+        if self.on_redraw_callback:
+            self.on_redraw_callback(right_click_event=event)
+
+    def set_redraw_callback(self, callback):
+        """Set callback function to trigger interaction logic."""
+        self.on_redraw_callback = callback
+
+    def get_image_coords(self, canvas_x, canvas_y, img_w, img_h):
+        """Converts canvas mouse coordinates to original image pixel coordinates."""
+        # Calculate center of canvas
+        cx = self.winfo_width() // 2 + self.pan_x
+        cy = self.winfo_height() // 2 + self.pan_y
+
+        # Calculate coordinates relative to the image center, accounting for zoom
+        rel_x = (canvas_x - cx) / self.zoom_level
+        rel_y = (canvas_y - cy) / self.zoom_level
+
+        # Translate from center-origin to top-left origin
+        img_x = int(rel_x + img_w / 2)
+        img_y = int(rel_y + img_h / 2)
+        return img_x, img_y
+
+    def _on_press(self, event):
+        self.last_mouse_x, self.last_mouse_y = event.x, event.y
+        if self.on_redraw_callback:
+            self.on_redraw_callback(click_event=event)
+
+    def _on_drag(self, event):
+        # Calculate movement delta in image pixels
+        dx = (event.x - self.last_mouse_x) / self.zoom_level
+        dy = (event.y - self.last_mouse_y) / self.zoom_level
+
+        # If no object is being edited, perform standard pan
+        if self.selected_obj_idx == -1 and not self.suspend_pan:
+            self.pan_x += event.x - self.last_mouse_x
+            self.pan_y += event.y - self.last_mouse_y
+
+        self.last_mouse_x, self.last_mouse_y = event.x, event.y
+
+        if self.on_redraw_callback:
+            # Main app handles moving points or boxes
+            self.on_redraw_callback(drag_event=(event, dx, dy))
+
+    def _on_release(self, event):
+        if self.on_redraw_callback:
+            self.on_redraw_callback()
+
+    def _on_motion(self, event):
+        if self.on_redraw_callback:
+            self.on_redraw_callback(move_event=event)
+
+    def _on_right_drag(self, event):
+        if self.on_redraw_callback:
+            self.on_redraw_callback(right_drag_event=event)
+
+    def _on_right_release(self, event):
+        if self.on_redraw_callback:
+            self.on_redraw_callback(right_release_event=event)
+
+    def display_image(self, cv_image):
+        """Display OpenCV image on canvas with current zoom and pan."""
+        if cv_image is None: return
+
+        # Pre-process for Tkinter
+        rgb_frame = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
+        pil_img = Image.fromarray(rgb_frame)
+
+        # Apply zoom resizing
+        new_w = int(pil_img.width * self.zoom_level)
+        new_h = int(pil_img.height * self.zoom_level)
+
+        if new_w > 0 and new_h > 0:
+            # Use NEAREST for maximum display throughput in live mode.
+            pil_img = pil_img.resize((new_w, new_h), Image.Resampling.NEAREST)
+
+        self.display_img = ImageTk.PhotoImage(pil_img)
+        self.delete("all")
+
+        # Center of canvas + current pan offset
+        cx = self.winfo_width() // 2 + self.pan_x
+        cy = self.winfo_height() // 2 + self.pan_y
+        self.create_image(cx, cy, image=self.display_img, anchor=tk.CENTER)
+
+    def _handle_zoom(self, event):
+        """Handle mouse wheel zoom centered on canvas."""
+        factor = 1.1 if (event.num == 4 or getattr(event, 'delta', 0) > 0) else 0.9
+        self.zoom_level = max(0.1, min(self.zoom_level * factor, 10.0))
+
+        # NEW: If the GUI has a zoom variable, keep it in sync
+        if hasattr(self.master.master, 'zoom_var'):
+            self.master.master.zoom_var.set(round(self.zoom_level, 2))
+
+        if self.on_redraw_callback:
+            self.on_redraw_callback()

src/gui/class_bank_mixin.py

@@ -0,0 +1,555 @@
+"""Class-set and reusable element-bank GUI mixin."""
+
+from pathlib import Path
+from datetime import datetime
+import json
+
+import numpy as np
+from tkinter import filedialog, messagebox, simpledialog
+
+
+class ClassBankMixin:
+    """Methods for class-set management and reusable element templates."""
+
+    def _bind_class_visibility_traces(self):
+        """Bind redraw callbacks for class visibility toggles once per class."""
+        for cls_name, var in self.ctrl_panel.visibility_vars.items():
+            if cls_name in self._visibility_trace_bound_classes:
+                continue
+            var.trace_add("write", lambda *args: self.redraw_current_frame())
+            self._visibility_trace_bound_classes.add(cls_name)
+
+        self._visibility_trace_bound_classes.intersection_update(
+            set(self.ctrl_panel.visibility_vars.keys())
+        )
+
+    def _refresh_class_ui(self):
+        """Refresh all class-dependent UI controls."""
+        class_names = list(self.class_params.keys())
+        self.ctrl_panel.update_class_selector(class_names)
+        self.ctrl_panel.update_class_visibility(class_names, self.class_colors)
+        self._bind_class_visibility_traces()
+
+        active = self.tk_vars['active_class'].get()
+        if active not in self.class_params:
+            self.tk_vars['active_class'].set(self._default_class_name())
+        self.redraw_current_frame()
+
+    def _apply_class_set(self, class_names, class_params=None, class_colors=None, annotation_mode=None):
+        """Apply a new class set and rebuild mappings/colors/params."""
+        cleaned = []
+        seen = set()
+        for raw in class_names:
+            name = str(raw).strip()
+            if not name or name in seen:
+                continue
+            seen.add(name)
+            cleaned.append(name)
+        if not cleaned:
+            raise ValueError("Class set is empty.")
+
+        old_params = dict(self.class_params)
+        old_colors = dict(self.class_colors)
+        class_params = class_params or {}
+        class_colors = class_colors or {}
+
+        self.class_params = {
+            name: dict(class_params.get(name, old_params.get(name, self._default_class_params())))
+            for name in cleaned
+        }
+        self.class_colors = {
+            name: class_colors.get(
+                name,
+                old_colors.get(name, self._class_color_palette[idx % len(self._class_color_palette)])
+            )
+            for idx, name in enumerate(cleaned)
+        }
+        self.class_to_id = {name: idx for idx, name in enumerate(cleaned)}
+        self.id_to_class = {idx: name for name, idx in self.class_to_id.items()}
+        if annotation_mode in self.annotation_mode_labels:
+            self.annotation_mode = annotation_mode
+        self.ctrl_panel.class_params = self.class_params
+        self._refresh_class_ui()
+
+    def choose_local_output_dir(self):
+        """Pick local folder used for manual image annotation saves."""
+        selected = filedialog.askdirectory(
+            title="Select local annotation output root",
+            initialdir=str(self.local_output_dir),
+        )
+        if not selected:
+            return
+        self.local_output_dir = Path(selected)
+        self.status_var.set(f"Local save folder: {self.local_output_dir}")
+
+    def create_new_class_set(self):
+        """Create a class set from a comma-separated list."""
+        mode = self._prompt_annotation_mode()
+        if mode is None:
+            return
+
+        current = ", ".join(self.class_to_id.keys())
+        entered = simpledialog.askstring(
+            "Create Class Set",
+            (
+                "Enter class names separated by commas.\n\n"
+                "Example: background, cell, nucleus\n\n"
+                f"Current classes: {current}"
+            ),
+            parent=self.root,
+        )
+        if entered is None:
+            return
+        try:
+            classes = [part.strip() for part in entered.split(",")]
+            self._apply_class_set(classes, annotation_mode=mode)
+            self.class_set_source_path = None
+            self.status_var.set(
+                f"Created class set with {len(self.class_to_id)} classes ({self._get_annotation_mode_label()})"
+            )
+        except Exception as e:
+            messagebox.showerror("Class Set Error", f"Failed to create class set:\n{e}")
+
+    def add_single_class(self):
+        """Add one class to current class set without replacing existing classes."""
+        entered = simpledialog.askstring(
+            "Add Class",
+            "Enter new class name:",
+            parent=self.root,
+        )
+        if entered is None:
+            return
+
+        new_name = entered.strip()
+        if not new_name:
+            messagebox.showwarning("Invalid Class", "Class name cannot be empty.")
+            return
+        if new_name in self.class_to_id:
+            messagebox.showwarning("Duplicate Class", f"Class '{new_name}' already exists.")
+            self.tk_vars['active_class'].set(new_name)
+            return
+
+        class_names = [name for _, name in sorted((idx, name) for name, idx in self.class_to_id.items())]
+        class_names.append(new_name)
+        self._apply_class_set(class_names, annotation_mode=self.annotation_mode)
+        self.tk_vars['active_class'].set(new_name)
+        self.status_var.set(f"Added class: {new_name} ({self._get_annotation_mode_label()})")
+
+    def load_class_set_from_disk(self):
+        """Load class set from JSON or TXT file."""
+        path = filedialog.askopenfilename(
+            title="Load class set",
+            filetypes=[
+                ("Class files", "*.json *.txt"),
+                ("JSON", "*.json"),
+                ("Text", "*.txt"),
+                ("All files", "*.*"),
+            ],
+        )
+        if not path:
+            return
+        try:
+            class_names, loaded_params, loaded_colors, loaded_mode = self._read_class_set(Path(path))
+            if loaded_mode is None:
+                loaded_mode = self._prompt_annotation_mode()
+                if loaded_mode is None:
+                    return
+            self._apply_class_set(class_names, loaded_params, loaded_colors, annotation_mode=loaded_mode)
+            self.class_set_source_path = Path(path)
+            self.status_var.set(
+                f"Loaded class set: {self.class_set_source_path.name} ({self._get_annotation_mode_label()})"
+            )
+        except Exception as e:
+            messagebox.showerror("Class Set Error", f"Failed loading class set:\n{e}")
+
+    def save_class_set_as(self):
+        """Persist current class set to JSON or TXT file."""
+        path = filedialog.asksaveasfilename(
+            title="Save class set as",
+            defaultextension=".json",
+            initialfile="class_set.json",
+            filetypes=[
+                ("JSON", "*.json"),
+                ("Text", "*.txt"),
+                ("All files", "*.*"),
+            ],
+        )
+        if not path:
+            return
+
+        out_path = Path(path)
+        class_names = [name for _, name in sorted((idx, name) for name, idx in self.class_to_id.items())]
+        try:
+            if out_path.suffix.lower() == ".txt":
+                out_path.write_text("\n".join(class_names) + ("\n" if class_names else ""), encoding="utf-8")
+            else:
+                payload = {
+                    "classes": class_names,
+                    "class_params": {name: self.class_params.get(name, self._default_class_params()) for name in class_names},
+                    "class_colors": {name: list(self.class_colors.get(name, (255, 0, 255))) for name in class_names},
+                    "annotation_mode": self.annotation_mode,
+                }
+                out_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+            self.class_set_source_path = out_path
+            self.status_var.set(f"Saved class set: {out_path.name}")
+        except Exception as e:
+            messagebox.showerror("Class Set Error", f"Failed saving class set:\n{e}")
+
+    def _read_class_set(self, path: Path):
+        """Read class names and optional style data from supported class-set files."""
+        suffix = path.suffix.lower()
+        if suffix == ".txt":
+            lines = path.read_text(encoding="utf-8").splitlines()
+            names = [line.strip() for line in lines if line.strip()]
+            return names, {}, {}, None
+
+        payload = json.loads(path.read_text(encoding="utf-8"))
+        if isinstance(payload, list):
+            names = [str(item).strip() for item in payload if str(item).strip()]
+            return names, {}, {}, None
+        if isinstance(payload, dict):
+            classes = payload.get("classes")
+            if isinstance(classes, list):
+                names = []
+                for item in classes:
+                    if isinstance(item, str):
+                        name = item.strip()
+                    elif isinstance(item, dict):
+                        name = str(item.get("name", "")).strip()
+                    else:
+                        name = str(item).strip()
+                    if name:
+                        names.append(name)
+                loaded_params = self._sanitize_loaded_class_params(payload.get("class_params", {}), names)
+                loaded_colors = self._sanitize_loaded_class_colors(payload.get("class_colors", {}), names)
+                loaded_mode = payload.get("annotation_mode", "yolo_single_polygon")
+                if loaded_mode not in self.annotation_mode_labels:
+                    loaded_mode = "yolo_single_polygon"
+                return names, loaded_params, loaded_colors, loaded_mode
+        raise ValueError("Expected TXT lines or JSON with a classes array.")
+
+    def _sanitize_loaded_class_params(self, maybe_params, class_names):
+        """Validate loaded class params map."""
+        if not isinstance(maybe_params, dict):
+            return {}
+        out = {}
+        for name in class_names:
+            raw = maybe_params.get(name)
+            if not isinstance(raw, dict):
+                continue
+            default = self._default_class_params()
+            out[name] = {
+                "low": int(raw.get("low", default["low"])),
+                "high": int(raw.get("high", default["high"])),
+                "clahe": float(raw.get("clahe", default["clahe"])),
+                "morph": int(raw.get("morph", default["morph"])),
+                "eps": float(raw.get("eps", default["eps"])),
+            }
+        return out
+
+    def _sanitize_loaded_class_colors(self, maybe_colors, class_names):
+        """Validate loaded class color map as BGR tuples."""
+        if not isinstance(maybe_colors, dict):
+            return {}
+        out = {}
+        for name in class_names:
+            raw = maybe_colors.get(name)
+            if not isinstance(raw, (list, tuple)) or len(raw) != 3:
+                continue
+            try:
+                bgr = tuple(max(0, min(255, int(v))) for v in raw)
+            except Exception:
+                continue
+            out[name] = bgr
+        return out
+
+    def _default_template_name(self, cls_name):
+        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+        return f"{cls_name}_{ts}"
+
+    @staticmethod
+    def _fit_box_inside_image(cx, cy, w, h, img_w, img_h):
+        w = max(2.0, min(float(w), float(img_w)))
+        h = max(2.0, min(float(h), float(img_h)))
+        x1 = float(cx) - w / 2.0
+        y1 = float(cy) - h / 2.0
+        x2 = x1 + w
+        y2 = y1 + h
+        if x1 < 0:
+            x2 -= x1
+            x1 = 0.0
+        if y1 < 0:
+            y2 -= y1
+            y1 = 0.0
+        if x2 > img_w:
+            shift = x2 - float(img_w)
+            x1 -= shift
+            x2 -= shift
+        if y2 > img_h:
+            shift = y2 - float(img_h)
+            y1 -= shift
+            y2 -= shift
+        x1 = max(0.0, x1)
+        y1 = max(0.0, y1)
+        x2 = min(float(img_w), x2)
+        y2 = min(float(img_h), y2)
+        return {"x1": x1, "y1": y1, "x2": x2, "y2": y2}
+
+    @staticmethod
+    def _normalize_poly(poly_arr, width, height):
+        arr = np.array(poly_arr, dtype=np.float32).reshape(-1, 2)
+        if len(arr) < 3:
+            return []
+        w = max(1.0, float(width))
+        h = max(1.0, float(height))
+        out = []
+        for x, y in arr:
+            nx = float(np.clip(x / w, 0.0, 1.0))
+            ny = float(np.clip(y / h, 0.0, 1.0))
+            out.append([round(nx, 6), round(ny, 6)])
+        return out
+
+    @staticmethod
+    def _denormalize_poly(poly_points, width, height):
+        pts = []
+        w = max(1.0, float(width))
+        h = max(1.0, float(height))
+        for pt in poly_points:
+            if not isinstance(pt, (list, tuple)) or len(pt) < 2:
+                continue
+            try:
+                nx = float(pt[0])
+                ny = float(pt[1])
+            except Exception:
+                continue
+            pts.append([np.clip(nx, 0.0, 1.0) * w, np.clip(ny, 0.0, 1.0) * h])
+        if len(pts) < 3:
+            return None
+        return np.array(pts, dtype=np.float32).reshape(-1, 1, 2)
+
+    def _element_bank_payload(self):
+        return {
+            "version": 1,
+            "updated_at": datetime.now().isoformat(),
+            "elements": self.element_bank,
+        }
+
+    def _load_default_element_bank_if_exists(self):
+        path = self.element_bank_path
+        if path.exists():
+            try:
+                self._load_element_bank_from_path(path, set_as_default=False)
+            except Exception as e:
+                self.status_var.set(f"Element bank load skipped: {e}")
+
+    def _load_element_bank_from_path(self, path: Path, set_as_default=True):
+        payload = json.loads(path.read_text(encoding="utf-8"))
+        elements = payload.get("elements", payload if isinstance(payload, list) else [])
+        if not isinstance(elements, list):
+            raise ValueError("Expected JSON object with elements[] or a top-level list.")
+        cleaned = []
+        for item in elements:
+            if not isinstance(item, dict):
+                continue
+            name = str(item.get("name", "")).strip()
+            class_name = str(item.get("class_name", self._default_class_name())).strip() or self._default_class_name()
+            try:
+                width = max(2.0, float(item.get("width", 64.0)))
+                height = max(2.0, float(item.get("height", 64.0)))
+            except Exception:
+                continue
+            polygons_norm = item.get("polygons_norm", [])
+            holes_norm = item.get("holes_norm", [])
+            if not isinstance(polygons_norm, list):
+                polygons_norm = []
+            if not isinstance(holes_norm, list):
+                holes_norm = []
+            if not name:
+                name = self._default_template_name(class_name)
+            cleaned.append(
+                {
+                    "name": name,
+                    "class_name": class_name,
+                    "width": width,
+                    "height": height,
+                    "polygons_norm": polygons_norm,
+                    "holes_norm": holes_norm,
+                    "created_at": str(item.get("created_at", datetime.now().isoformat())),
+                }
+            )
+        self.element_bank = cleaned
+        if set_as_default:
+            self.element_bank_path = path
+        self.status_var.set(f"Loaded element bank: {path.name} ({len(self.element_bank)} elements)")
+
+    def load_element_bank_from_disk(self):
+        path = filedialog.askopenfilename(
+            title="Load element bank",
+            filetypes=[("JSON", "*.json"), ("All files", "*.*")],
+        )
+        if not path:
+            return
+        try:
+            self._load_element_bank_from_path(Path(path), set_as_default=True)
+        except Exception as e:
+            messagebox.showerror("Element Bank Error", f"Failed loading element bank:\n{e}")
+
+    def _save_element_bank_to_path(self, path: Path, set_as_default=True):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(self._element_bank_payload(), indent=2), encoding="utf-8")
+        if set_as_default:
+            self.element_bank_path = path
+        self.status_var.set(f"Saved element bank: {path.name}")
+
+    def save_element_bank_as(self):
+        path = filedialog.asksaveasfilename(
+            title="Save element bank as",
+            defaultextension=".json",
+            initialfile="element_bank.json",
+            filetypes=[("JSON", "*.json"), ("All files", "*.*")],
+        )
+        if not path:
+            return
+        try:
+            self._save_element_bank_to_path(Path(path), set_as_default=True)
+        except Exception as e:
+            messagebox.showerror("Element Bank Error", f"Failed saving element bank:\n{e}")
+
+    def _pick_element_from_bank(self):
+        if not self.element_bank:
+            messagebox.showinfo(
+                "Element Bank",
+                "Element bank is empty.\nSelect an annotation and use Elements -> Add Selected to Bank.",
+            )
+            return None
+        options = [
+            f"{idx + 1}. {el.get('name', 'unnamed')} [{el.get('class_name', '?')}]"
+            for idx, el in enumerate(self.element_bank)
+        ]
+        prompt = "Choose element number:\n\n" + "\n".join(options[:30])
+        if len(options) > 30:
+            prompt += f"\n... ({len(options) - 30} more not shown)"
+        picked = simpledialog.askinteger("Place Element", prompt, parent=self.root, minvalue=1, maxvalue=len(options))
+        if not picked:
+            return None
+        return self.element_bank[picked - 1]
+
+    def add_selected_to_element_bank(self):
+        idx = self.canvas.selected_obj_idx
+        if idx == -1 or idx >= len(self.last_results):
+            messagebox.showwarning("Element Bank", "Select an annotation first.")
+            return
+        det = self.last_results[idx]
+        box = det.get("box", det)
+        w = max(2.0, float(box.get("x2", 0.0)) - float(box.get("x1", 0.0)))
+        h = max(2.0, float(box.get("y2", 0.0)) - float(box.get("y1", 0.0)))
+        class_name = str(det.get("name", det.get("label", self._default_class_name())))
+        polygons_rel = self._det_get_polygons_relative(det)
+        holes_rel = self._det_get_holes_relative(det)
+        if not polygons_rel:
+            polygons_rel = [np.array([[[0.0, 0.0]], [[w, 0.0]], [[w, h]], [[0.0, h]]], dtype=np.float32)]
+        polygons_norm = [self._normalize_poly(poly, w, h) for poly in polygons_rel]
+        holes_norm = [self._normalize_poly(poly, w, h) for poly in holes_rel]
+        name = simpledialog.askstring(
+            "Element Name",
+            "Name for this reusable element:",
+            initialvalue=self._default_template_name(class_name),
+            parent=self.root,
+        )
+        if name is None:
+            return
+        name = name.strip()
+        if not name:
+            messagebox.showwarning("Element Bank", "Element name cannot be empty.")
+            return
+        self.element_bank.append(
+            {
+                "name": name,
+                "class_name": class_name,
+                "width": w,
+                "height": h,
+                "polygons_norm": polygons_norm,
+                "holes_norm": holes_norm,
+                "created_at": datetime.now().isoformat(),
+            }
+        )
+        try:
+            self._save_element_bank_to_path(self.element_bank_path, set_as_default=True)
+        except Exception as e:
+            messagebox.showwarning("Element Bank", f"Added to memory, but auto-save failed:\n{e}")
+        self.status_var.set(f"Added element: {name}")
+
+    def place_element_from_bank(self):
+        if self.current_raw_frame is None:
+            messagebox.showwarning("Warning", "No image loaded.")
+            return
+        element = self._pick_element_from_bank()
+        if element is None:
+            return
+
+        scale = simpledialog.askfloat(
+            "Element Scale",
+            "Scale factor (1.0 = original size):",
+            initialvalue=1.0,
+            minvalue=0.1,
+            maxvalue=10.0,
+            parent=self.root,
+        )
+        if scale is None:
+            return
+
+        img_h, img_w = self.current_raw_frame.shape[:2]
+        cw = self.canvas.winfo_width() / 2
+        ch = self.canvas.winfo_height() / 2
+        cx, cy = self.canvas.get_image_coords(cw, ch, img_w, img_h)
+        width = max(2.0, float(element.get("width", 64.0)) * float(scale))
+        height = max(2.0, float(element.get("height", 64.0)) * float(scale))
+        box = self._fit_box_inside_image(cx, cy, width, height, img_w, img_h)
+        box_w = max(1.0, box["x2"] - box["x1"])
+        box_h = max(1.0, box["y2"] - box["y1"])
+
+        polygons = []
+        for poly in element.get("polygons_norm", []):
+            arr = self._denormalize_poly(poly, box_w, box_h)
+            if arr is not None:
+                polygons.append(arr)
+        if not polygons:
+            polygons = [np.array([[[0.0, 0.0]], [[box_w, 0.0]], [[box_w, box_h]], [[0.0, box_h]]], dtype=np.float32)]
+
+        holes = []
+        for poly in element.get("holes_norm", []):
+            arr = self._denormalize_poly(poly, box_w, box_h)
+            if arr is not None:
+                holes.append(arr)
+
+        class_name = str(element.get("class_name", self._default_class_name()))
+        if class_name not in self.class_to_id:
+            if messagebox.askyesno(
+                "Unknown Class",
+                f"Element class '{class_name}' is not in current class set.\nAdd it now?",
+                parent=self.root,
+            ):
+                class_names = [name for _, name in sorted((idx, name) for name, idx in self.class_to_id.items())]
+                class_names.append(class_name)
+                self._apply_class_set(class_names, annotation_mode=self.annotation_mode)
+            else:
+                class_name = self.tk_vars['active_class'].get()
+
+        self._save_history()
+        det = {
+            "box": box,
+            "name": class_name,
+            "poly": polygons[0],
+            "polygons": polygons,
+            "conf": 1.0,
+        }
+        if holes:
+            det["holes"] = holes
+
+        self.last_results.append(det)
+        self.canvas.selected_obj_idx = len(self.last_results) - 1
+        self.canvas.selected_point_idx = -1
+        self.selected_island_idx = 0
+        self.selected_island_kind = "outer"
+        self.redraw_current_frame()
+        self.status_var.set(f"Placed element: {element.get('name', 'unnamed')}")

src/gui/control_panel.py

@@ -0,0 +1,337 @@
+"""Control panel GUI component with parameter controls and focus plot."""
+import tkinter as tk
+from tkinter import ttk
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
+
+
+class ControlPanel(tk.Frame):
+    """Right-side control panel with tuning parameters and focus plot."""
+
+    def __init__(self, parent, class_params, tk_vars, callbacks):
+        """
+        Initialize control panel.
+
+        Args:
+            parent: Parent tkinter widget
+            class_params: Dict of class-specific parameters
+            tk_vars: Dict with tkinter variables:
+                - active_class: StringVar for selected class
+                - canny_high: IntVar for Canny high threshold
+                - morph_kernel: IntVar for morphological kernel size
+                - poly_epsilon: DoubleVar for polygon epsilon
+                - show_debug_edges: BooleanVar for edge visualization
+                - plot_window_size: IntVar for focus plot window
+                - annotation_line_width: IntVar for annotation line width
+            callbacks: Dict with callback functions:
+                - on_train_click: Function to save for training
+        """
+        super().__init__(parent, width=300, relief=tk.SUNKEN, bd=1)
+        self.class_params = class_params
+        self.tk_vars = tk_vars
+        self.callbacks = callbacks
+
+        # Create canvas and scrollbar for scrollable content
+        self.canvas = tk.Canvas(self, width=280, highlightthickness=0)
+        self.scrollbar = tk.Scrollbar(self, orient="vertical", command=self.canvas.yview)
+        self.scrollable_frame = tk.Frame(self.canvas)
+
+        self.scrollable_frame.bind(
+            "<Configure>",
+            lambda e: self.canvas.configure(scrollregion=self.canvas.bbox("all"))
+        )
+
+        self.canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
+        self.canvas.configure(yscrollcommand=self.scrollbar.set)
+
+        self.canvas.pack(side="left", fill="both", expand=True, padx=5, pady=5)
+        self.scrollbar.pack(side="right", fill="y")
+
+        # FIX: Only bind to this specific canvas, not bind_all
+        self.canvas.bind("<MouseWheel>", self._on_mousewheel)
+        self.canvas.bind("<Button-4>", self._on_mousewheel)
+        self.canvas.bind("<Button-5>", self._on_mousewheel)
+
+        self._setup_ui()
+
+    def _on_mousewheel(self, event):
+        """Handle mousewheel scrolling."""
+        if event.num == 4 or event.delta > 0:
+            self.canvas.yview_scroll(-1, "units")
+        elif event.num == 5 or event.delta < 0:
+            self.canvas.yview_scroll(1, "units")
+
+    def _setup_ui(self):
+        """Setup control panel UI elements."""
+        frame = self.scrollable_frame  # Use scrollable frame for all widgets
+
+        # Title
+        tk.Label(frame, text="Tuning Controls", font=('bold', 12)).pack(anchor=tk.W, pady=5)
+
+        # Class selector
+        tk.Label(frame, text="Active Class:").pack(anchor=tk.W)
+        self.class_combobox = ttk.Combobox(
+            frame,
+            textvariable=self.tk_vars['active_class'],
+            values=list(self.class_params.keys()),
+            state="readonly"
+        )
+        self.class_combobox.pack(fill=tk.X, pady=5)
+
+        # Canny High threshold
+        tk.Label(frame, text="Canny High:").pack(anchor=tk.W)
+        tk.Scale(
+            frame,
+            from_=0,
+            to=255,
+            orient=tk.HORIZONTAL,
+            variable=self.tk_vars['canny_high']
+        ).pack(fill=tk.X)
+
+        # Morph kernel size
+        tk.Label(frame, text="Morph Kernel:").pack(anchor=tk.W)
+        tk.Scale(
+            frame,
+            from_=1,
+            to=100,
+            orient=tk.HORIZONTAL,
+            variable=self.tk_vars['morph_kernel']
+        ).pack(fill=tk.X)
+
+        # Polygon precision
+        tk.Label(frame, text="Poly Precision:").pack(anchor=tk.W)
+        tk.Scale(
+            frame,
+            from_=0.0001,
+            to=0.05,
+            resolution=0.0001,
+            orient=tk.HORIZONTAL,
+            variable=self.tk_vars['poly_epsilon']
+        ).pack(fill=tk.X)
+
+        # Debug edges checkbox
+        tk.Checkbutton(
+            frame,
+            text="View Raw Edge Map",
+            variable=self.tk_vars['show_debug_edges']
+        ).pack(anchor=tk.W, pady=5)
+
+        # Segmentation method selector
+        tk.Label(frame, text="Segmentation Method:", font=('bold', 10)).pack(anchor=tk.W, pady=(10, 5))
+
+        seg_frame = tk.Frame(frame)
+        seg_frame.pack(fill=tk.X, pady=5)
+
+        # Ensure we use the correct variable from tk_vars
+        seg_var = self.tk_vars.get('segmentation_method')
+
+        tk.Radiobutton(
+            seg_frame,
+            text="Watershed",
+            variable=seg_var,
+            value='watershed'
+        ).pack(side=tk.LEFT, padx=2)
+
+        tk.Radiobutton(
+            seg_frame,
+            text="GrabCut",
+            variable=seg_var,
+            value='grabcut'
+        ).pack(side=tk.LEFT, padx=2)
+
+        tk.Radiobutton(
+            seg_frame,
+            text="✨ Smart",
+            variable=seg_var,
+            value='smart'
+        ).pack(side=tk.LEFT, padx=2)
+
+        # Brush tool section
+        tk.Label(frame, text="Brush Tool:", font=('bold', 10)).pack(anchor=tk.W, pady=(10, 5))
+
+        # Brush usage hint
+        hint_label = tk.Label(
+            frame,
+            text="💡 Tip: Paint BACKGROUND (red)\naround object edges for best results",
+            font=('Arial', 9, 'italic'),
+            fg='#666',
+            justify=tk.LEFT
+        )
+        hint_label.pack(anchor=tk.W, pady=2)
+
+        # Brush mode toggle
+        brush_frame = tk.Frame(frame)
+        brush_frame.pack(fill=tk.X, pady=5)
+
+        tk.Checkbutton(
+            brush_frame,
+            text="🖌️ Enable Brush Mode",
+            variable=self.tk_vars['brush_mode']
+        ).pack(side=tk.LEFT)
+
+        # Brush type radio buttons
+        brush_type_frame = tk.Frame(frame)
+        brush_type_frame.pack(fill=tk.X, pady=2)
+
+        tk.Radiobutton(
+            brush_type_frame,
+            text="Background",
+            variable=self.tk_vars['brush_type'],
+            value='background'
+        ).pack(side=tk.LEFT, padx=5)
+
+        tk.Radiobutton(
+            brush_type_frame,
+            text="Foreground",
+            variable=self.tk_vars['brush_type'],
+            value='foreground'
+        ).pack(side=tk.LEFT, padx=5)
+
+        # Brush size slider
+        tk.Label(frame, text="Brush Size:").pack(anchor=tk.W)
+        tk.Scale(
+            frame,
+            from_=5,
+            to=50,
+            orient=tk.HORIZONTAL,
+            variable=self.tk_vars['brush_size']
+        ).pack(fill=tk.X)
+
+        # Regenerate edge button
+        tk.Button(
+            frame,
+            text="🔄 Regenerate Edge",
+            command=self.callbacks.get('on_regenerate_edge', lambda: None),
+            bg="#007bff",
+            fg="white"
+        ).pack(fill=tk.X, pady=5)
+
+        # Clear brush strokes button
+        tk.Button(
+            frame,
+            text="Clear Brush Strokes",
+            command=self.callbacks.get('on_clear_brush', lambda: None)
+        ).pack(fill=tk.X, pady=2)
+
+        # Class visibility section
+        tk.Label(frame, text="Class Visibility:", font=('bold', 10)).pack(anchor=tk.W, pady=(10, 5))
+        self.class_visibility_frame = tk.Frame(frame)
+        self.class_visibility_frame.pack(fill=tk.X, pady=5)
+
+        # Store visibility variables (will be populated by update_class_visibility)
+        self.visibility_vars = {}
+
+        # Annotation stroke width
+        tk.Label(frame, text="Annotation Line Width:").pack(anchor=tk.W, pady=(8, 0))
+        tk.Scale(
+            frame,
+            from_=1,
+            to=6,
+            orient=tk.HORIZONTAL,
+            variable=self.tk_vars['annotation_line_width']
+        ).pack(fill=tk.X)
+
+        # Training button
+        tk.Button(
+            frame,
+            text="SAVE FOR TRAINING",
+            bg="green",
+            fg="white",
+            font=('bold'),
+            command=self.callbacks['on_train_click']
+        ).pack(fill=tk.X, pady=20)
+
+        # Focus plot section
+        self._setup_focus_plot()
+
+    def _setup_focus_plot(self):
+        """Setup focus history plot."""
+        frame = self.scrollable_frame
+
+        tk.Label(frame, text="Focus History", font=('bold')).pack(anchor=tk.W, pady=(10, 0))
+
+        # Window size slider
+        tk.Scale(
+            frame,
+            from_=50,
+            to=500,
+            orient=tk.HORIZONTAL,
+            variable=self.tk_vars['plot_window_size']
+        ).pack(fill=tk.X)
+
+        # Create matplotlib figure
+        self.fig, self.ax = plt.subplots(figsize=(3, 2), dpi=80)
+        self.ax.set_facecolor('#f0f0f0')
+        self.line, = self.ax.plot([], [], color='blue', lw=1.5)
+        self.ax.get_xaxis().set_visible(False)
+
+        # Embed in tkinter
+        self.canvas_plot = FigureCanvasTkAgg(self.fig, master=frame)
+        self.canvas_plot.get_tk_widget().pack(fill=tk.X, pady=10)
+
+    def update_focus_plot(self, focus_data, window_size):
+        """
+        Update focus plot with new data.
+
+        Args:
+            focus_data: List of focus scores
+            window_size: Number of recent points to display
+        """
+        display_data = focus_data[-window_size:] if len(focus_data) > window_size else focus_data
+
+        if not display_data:
+            self.line.set_data([], [])
+        else:
+            self.line.set_data(range(len(display_data)), display_data)
+            ymin, ymax = min(display_data), max(display_data)
+            margin = (ymax - ymin) * 0.1 if ymax > ymin else 1.0
+            self.ax.set_ylim(ymin - margin, ymax + margin)
+            self.ax.set_xlim(0, window_size)
+
+        self.canvas_plot.draw_idle()
+
+    def update_class_visibility(self, class_names, class_colors):
+        """
+        Update class visibility checkboxes.
+
+        Args:
+            class_names: List of class names
+            class_colors: Dict mapping class names to BGR colors
+        """
+        class_set = set(class_names)
+
+        # Remove stale variables for classes no longer present
+        for existing in list(self.visibility_vars.keys()):
+            if existing not in class_set:
+                del self.visibility_vars[existing]
+
+        # Clear existing checkboxes
+        for widget in self.class_visibility_frame.winfo_children():
+            widget.destroy()
+
+        # Create checkbox for each class
+        for cls_name in class_names:
+            if cls_name not in self.visibility_vars:
+                self.visibility_vars[cls_name] = tk.BooleanVar(value=True)
+
+            frame = tk.Frame(self.class_visibility_frame)
+            frame.pack(fill=tk.X, pady=2)
+
+            # Color indicator
+            color_bgr = class_colors.get(cls_name, (255, 0, 255))
+            # Convert BGR to hex for tkinter
+            color_hex = f"#{color_bgr[2]:02x}{color_bgr[1]:02x}{color_bgr[0]:02x}"
+            color_label = tk.Label(frame, bg=color_hex, width=2, height=1)
+            color_label.pack(side=tk.LEFT, padx=(0, 5))
+
+            # Checkbox
+            cb = tk.Checkbutton(
+                frame,
+                text=cls_name,
+                variable=self.visibility_vars[cls_name]
+            )
+            cb.pack(side=tk.LEFT)
+
+    def update_class_selector(self, class_names):
+        """Refresh active-class dropdown with a new class set."""
+        self.class_combobox.configure(values=list(class_names))

src/gui/dataset_prep_mixin.py

@@ -0,0 +1,522 @@
+"""YOLO dataset prep script generation mixin for the GUI."""
+
+from pathlib import Path
+from datetime import datetime
+import json
+import os
+import textwrap
+
+from tkinter import filedialog, messagebox, simpledialog
+
+
+class DatasetPrepMixin:
+    """Methods to generate standalone YOLO segmentation dataset prep scripts."""
+
+    def show_dataset_prep_script_dialog(self):
+        """Collect settings and generate a headless dataset preparation script."""
+        source_dir = filedialog.askdirectory(
+            title="Select source folder (contains images/labels or images+labels folders)"
+        )
+        if not source_dir:
+            return
+
+        default_dest = str(Path(source_dir).parent / f"{Path(source_dir).name}_yolo_seg")
+        destination_dir = simpledialog.askstring(
+            "Destination Dataset Path",
+            (
+                "Destination path to create on the training machine.\n"
+                "This can be an absolute path and does not need to exist yet."
+            ),
+            initialvalue=default_dest,
+            parent=self.root,
+        )
+        if destination_dir is None:
+            return
+        destination_dir = destination_dir.strip()
+        if not destination_dir:
+            messagebox.showwarning("Invalid Input", "Destination path cannot be empty.")
+            return
+
+        train_ratio = simpledialog.askfloat(
+            "Training Ratio",
+            "Training split ratio (e.g. 0.85 for 85/15 train/val):",
+            initialvalue=0.85,
+            minvalue=0.5,
+            maxvalue=0.99,
+            parent=self.root,
+        )
+        if train_ratio is None:
+            return
+
+        seed = simpledialog.askinteger(
+            "Random Seed",
+            "Random seed for reproducible split:",
+            initialvalue=42,
+            parent=self.root,
+        )
+        if seed is None:
+            return
+
+        model_spec = self._prompt_training_model_spec()
+        if model_spec is None:
+            return
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        default_script_name = f"prepare_yolo_seg_dataset_{timestamp}.py"
+        script_path = filedialog.asksaveasfilename(
+            title="Save generated preparation script",
+            defaultextension=".py",
+            initialfile=default_script_name,
+            filetypes=[("Python Script", "*.py"), ("All Files", "*.*")],
+        )
+        if not script_path:
+            return
+
+        class_names = [
+            name for _, name in sorted((class_id, name) for name, class_id in self.class_to_id.items())
+        ]
+
+        try:
+            project_root = Path(__file__).resolve().parents[2]
+            source_for_script = self._to_project_relative(source_dir, project_root)
+            destination_for_script = self._to_project_relative(destination_dir, project_root)
+            script_content = self._build_yolo_prep_script(
+                source_dir=source_for_script,
+                destination_dir=destination_for_script,
+                train_ratio=train_ratio,
+                seed=seed,
+                class_names=class_names,
+                model_spec=model_spec,
+            )
+            output_path = Path(script_path)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            output_path.write_text(script_content, encoding="utf-8")
+            os.chmod(output_path, 0o755)
+        except Exception as e:
+            messagebox.showerror("Generation Error", f"Failed to generate script:\n{e}")
+            self.status_var.set("Dataset prep script generation failed")
+            return
+
+        messagebox.showinfo(
+            "Script Generated",
+            (
+                f"Script saved:\n{output_path}\n\n"
+                "Run this on the training machine:\n"
+                f"python3 {output_path}\n\n"
+                "Optional dry-run:\n"
+                f"python3 {output_path} --dry-run"
+            ),
+        )
+        self.status_var.set(f"Generated dataset prep script: {output_path.name}")
+
+    def _prompt_training_model_spec(self):
+        """Prompt for initial training model: local file in models/ or Ultralytics model name."""
+        use_local = messagebox.askyesnocancel(
+            "Training Model",
+            (
+                "Use a local model file from this project (recommended if available)?\n\n"
+                "Yes: choose a local .pt/.yaml model file\n"
+                "No: enter an Ultralytics model name (e.g. yolo11n-seg.pt)\n"
+                "Cancel: abort"
+            ),
+            parent=self.root,
+        )
+
+        if use_local is None:
+            return None
+
+        if use_local:
+            project_root = Path(__file__).resolve().parents[2]
+            initial_dir = project_root / "models"
+            model_path = filedialog.askopenfilename(
+                title="Select input model",
+                initialdir=str(initial_dir if initial_dir.exists() else project_root),
+                filetypes=[
+                    ("YOLO models/config", "*.pt *.yaml *.yml"),
+                    ("All Files", "*.*"),
+                ],
+            )
+            if not model_path:
+                return None
+            return self._to_project_relative(model_path, project_root)
+
+        model_name = simpledialog.askstring(
+            "Ultralytics Model Name",
+            (
+                "Enter Ultralytics model name to auto-download at training time.\n"
+                "Examples: yolo11n-seg.pt, yolo11s-seg.pt"
+            ),
+            initialvalue="yolo11n-seg.pt",
+            parent=self.root,
+        )
+        if model_name is None:
+            return None
+        model_name = model_name.strip()
+        if not model_name:
+            messagebox.showwarning("Invalid Input", "Model name cannot be empty.")
+            return None
+        return model_name
+
+    @staticmethod
+    def _to_project_relative(path_value, project_root: Path):
+        """Convert absolute path under project root to relative POSIX path for portability."""
+        raw = Path(path_value).expanduser()
+        if not raw.is_absolute():
+            return raw.as_posix()
+        try:
+            return raw.resolve().relative_to(project_root.resolve()).as_posix()
+        except ValueError:
+            return str(raw.resolve())
+
+    def _build_yolo_prep_script(self, source_dir, destination_dir, train_ratio, seed, class_names, model_spec):
+        """Build a standalone Python script to prepare YOLO segmentation dataset."""
+        generated_at = datetime.now().isoformat(timespec="seconds")
+        script = f'''#!/usr/bin/env python3
+"""
+Auto-generated by AareLC ML Studio on {generated_at}.
+Prepare YOLO segmentation train/val folders from source images/labels.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import shutil
+from collections import Counter
+from datetime import datetime, timezone
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+DEFAULT_SOURCE = {source_dir!r}
+DEFAULT_DESTINATION = {destination_dir!r}
+DEFAULT_TRAIN_RATIO = {train_ratio!r}
+DEFAULT_SEED = {seed!r}
+DEFAULT_MODEL_SPEC = {model_spec!r}
+CLASS_NAMES = {json.dumps(class_names)}
+IMAGE_EXTENSIONS = {{".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}}
+
+
+def resolve_user_path(path_value: str) -> Path:
+    p = Path(path_value).expanduser()
+    if not p.is_absolute():
+        p = PROJECT_ROOT / p
+    return p.resolve()
+
+
+def resolve_layout(source_root: Path) -> tuple[Path, Path]:
+    images_dir = source_root / "images"
+    labels_dir = source_root / "labels"
+    if images_dir.is_dir() and labels_dir.is_dir():
+        return images_dir, labels_dir
+    return source_root, source_root
+
+
+def collect_pairs(source_root: Path):
+    images_dir, labels_dir = resolve_layout(source_root)
+    image_files = []
+    for ext in IMAGE_EXTENSIONS:
+        image_files.extend(images_dir.rglob(f"*{{ext}}"))
+        image_files.extend(images_dir.rglob(f"*{{ext.upper()}}"))
+    image_files = sorted(set(image_files))
+
+    pairs = []
+    missing_labels = []
+    used_labels = set()
+
+    for img_path in image_files:
+        rel_image = img_path.relative_to(images_dir)
+        label_candidate = (labels_dir / rel_image).with_suffix(".txt")
+        if not label_candidate.exists():
+            flat_fallback = labels_dir / f"{{img_path.stem}}.txt"
+            if flat_fallback.exists():
+                label_candidate = flat_fallback
+
+        if label_candidate.exists():
+            pairs.append((img_path, label_candidate, rel_image))
+            used_labels.add(label_candidate.resolve())
+        else:
+            missing_labels.append(str(img_path))
+
+    all_labels = {{p.resolve() for p in labels_dir.rglob("*.txt")}}
+    orphan_labels = sorted(str(p) for p in (all_labels - used_labels))
+    return pairs, missing_labels, orphan_labels, images_dir, labels_dir
+
+
+def split_pairs(pairs, train_ratio: float, seed: int):
+    ordered = list(pairs)
+    rng = random.Random(seed)
+    rng.shuffle(ordered)
+    total = len(ordered)
+    if total == 0:
+        return [], []
+
+    train_count = int(total * train_ratio)
+    if total > 1:
+        train_count = max(1, min(total - 1, train_count))
+    else:
+        train_count = 1
+    return ordered[:train_count], ordered[train_count:]
+
+
+def parse_class_ids(label_path: Path):
+    class_ids = []
+    for line in label_path.read_text(encoding="utf-8").splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        head = line.split()[0]
+        try:
+            class_ids.append(int(float(head)))
+        except ValueError:
+            continue
+    return class_ids
+
+
+def write_dataset_yaml(destination_root: Path):
+    lines = [
+        f"path: {{destination_root.resolve()}}",
+        "train: train/images",
+        "val: val/images",
+        "",
+        "names:",
+    ]
+    for class_id, class_name in enumerate(CLASS_NAMES):
+        lines.append(f"  {{class_id}}: {{class_name}}")
+    dataset_yaml = destination_root / "dataset.yaml"
+    dataset_yaml.write_text("\\n".join(lines) + "\\n", encoding="utf-8")
+    return dataset_yaml
+
+
+def write_train_script(destination_root: Path, dataset_yaml: Path, model_spec: str):
+    train_script = destination_root / "train_yolo_seg.py"
+    script_text = f"""#!/usr/bin/env python3
+from pathlib import Path
+from ultralytics import YOLO
+import argparse
+
+
+def resolve_model(model_value: str) -> str:
+    p = Path(model_value).expanduser()
+    if p.is_absolute():
+        return str(p)
+    local_candidate = Path(__file__).resolve().parent.parent / p
+    if local_candidate.exists():
+        return str(local_candidate.resolve())
+    return model_value
+
+
+def main():
+    parser = argparse.ArgumentParser(description=\\"Train YOLO segmentation model\\")
+    parser.add_argument(\\"--model\\", default={model_spec!r}, help=\\"Local model path or Ultralytics model name\\")
+    parser.add_argument(\\"--data\\", default=None, help=\\"Dataset YAML path (default: ./dataset.yaml next to this script)\\")
+    parser.add_argument(\\"--epochs\\", type=int, default=100)
+    parser.add_argument(\\"--imgsz\\", type=int, default=640)
+    parser.add_argument(\\"--batch\\", type=int, default=16)
+    parser.add_argument(\\"--device\\", default=\\"0\\")
+    parser.add_argument(\\"--project\\", default=\\"runs/segment\\")
+    parser.add_argument(\\"--name\\", default=\\"aarelc_seg_train\\")
+    args = parser.parse_args()
+
+    model_value = resolve_model(args.model)
+    if args.data:
+        data_yaml = Path(args.data).expanduser().resolve()
+    else:
+        data_yaml = Path(__file__).resolve().parent / "dataset.yaml"
+    if not data_yaml.exists():
+        raise FileNotFoundError(
+            f"Dataset YAML not found: {{data_yaml}}. "
+            "Run prepare_yolo_seg_dataset first or pass --data /path/to/dataset.yaml."
+        )
+
+    model = YOLO(model_value)
+    model.train(
+        data=str(data_yaml),
+        epochs=args.epochs,
+        patience=50,
+        amp=False,
+        overlap_mask=False,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        device=args.device,
+        project=args.project,
+        name=args.name,
+    )
+
+
+if __name__ == \\"__main__\\":
+    main()
+"""
+    train_script.write_text(script_text, encoding="utf-8")
+    train_script.chmod(0o755)
+    return train_script
+
+
+def copy_pairs(split_name: str, rows, destination_root: Path, dry_run: bool):
+    log_rows = []
+    for image_path, label_path, rel_image in rows:
+        dst_image = destination_root / split_name / "images" / rel_image
+        dst_label = destination_root / split_name / "labels" / rel_image.with_suffix(".txt")
+        log_rows.append(
+            {{
+                "source_image": str(image_path),
+                "source_label": str(label_path),
+                "dest_image": str(dst_image),
+                "dest_label": str(dst_label),
+            }}
+        )
+        if dry_run:
+            continue
+        dst_image.parent.mkdir(parents=True, exist_ok=True)
+        dst_label.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(image_path, dst_image)
+        shutil.copy2(label_path, dst_label)
+    return log_rows
+
+
+def build_text_log(metadata: dict) -> str:
+    lines = [
+        "YOLO Dataset Preparation Log",
+        "=" * 32,
+        f"Prepared at (UTC): {{metadata['prepared_at_utc']}}",
+        f"Source root: {{metadata['source_root']}}",
+        f"Destination root: {{metadata['destination_root']}}",
+        f"Model spec: {{metadata['model_spec']}}",
+        f"Training script: {{metadata['train_script']}}",
+        f"Dry run: {{metadata['dry_run']}}",
+        f"Train ratio: {{metadata['train_ratio']}}",
+        f"Seed: {{metadata['seed']}}",
+        "",
+        f"Total image+label pairs: {{metadata['totals']['pairs']}}",
+        f"Train pairs: {{metadata['totals']['train_pairs']}}",
+        f"Validation pairs: {{metadata['totals']['val_pairs']}}",
+        f"Missing labels: {{metadata['totals']['missing_labels']}}",
+        f"Orphan labels: {{metadata['totals']['orphan_labels']}}",
+        "",
+        "Class counts:",
+    ]
+    for class_id, count in sorted(metadata["class_counts"].items(), key=lambda x: int(x[0])):
+        class_name = metadata["class_names"].get(class_id, f"class_{{class_id}}")
+        lines.append(f"  class {{class_id}} ({{class_name}}): {{count}}")
+
+    lines.append("")
+    lines.append("Train files:")
+    for row in metadata["train_files"]:
+        lines.append(f"  IMG {{row['source_image']}} -> {{row['dest_image']}}")
+        lines.append(f"  LBL {{row['source_label']}} -> {{row['dest_label']}}")
+
+    lines.append("")
+    lines.append("Validation files:")
+    for row in metadata["val_files"]:
+        lines.append(f"  IMG {{row['source_image']}} -> {{row['dest_image']}}")
+        lines.append(f"  LBL {{row['source_label']}} -> {{row['dest_label']}}")
+
+    if metadata["missing_labels"]:
+        lines.append("")
+        lines.append("Images missing labels:")
+        for path in metadata["missing_labels"]:
+            lines.append(f"  {{path}}")
+
+    if metadata["orphan_labels"]:
+        lines.append("")
+        lines.append("Labels without matching images:")
+        for path in metadata["orphan_labels"]:
+            lines.append(f"  {{path}}")
+
+    return "\\n".join(lines) + "\\n"
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare YOLO segmentation train/val dataset.")
+    parser.add_argument("--source", default=DEFAULT_SOURCE, help="Source folder with images/labels")
+    parser.add_argument("--destination", default=DEFAULT_DESTINATION, help="Destination dataset root")
+    parser.add_argument("--train-ratio", type=float, default=DEFAULT_TRAIN_RATIO, help="Train split ratio")
+    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="Random seed")
+    parser.add_argument("--dry-run", action="store_true", help="Preview only, do not copy files")
+    args = parser.parse_args()
+
+    if not (0.0 < args.train_ratio < 1.0):
+        raise ValueError("train-ratio must be between 0 and 1")
+
+    source_root = resolve_user_path(args.source)
+    destination_root = resolve_user_path(args.destination)
+    if not source_root.exists():
+        raise FileNotFoundError(f"Source folder not found: {{source_root}}")
+
+    pairs, missing_labels, orphan_labels, images_dir, labels_dir = collect_pairs(source_root)
+    train_pairs, val_pairs = split_pairs(pairs, args.train_ratio, args.seed)
+
+    class_counts = Counter()
+    for _, label_path, _ in pairs:
+        class_counts.update(parse_class_ids(label_path))
+
+    if not args.dry_run:
+        for split in ("train", "val"):
+            (destination_root / split / "images").mkdir(parents=True, exist_ok=True)
+            (destination_root / split / "labels").mkdir(parents=True, exist_ok=True)
+
+    train_rows = copy_pairs("train", train_pairs, destination_root, args.dry_run)
+    val_rows = copy_pairs("val", val_pairs, destination_root, args.dry_run)
+
+    dataset_yaml = None
+    train_script_path = None
+    if not args.dry_run:
+        dataset_yaml = write_dataset_yaml(destination_root)
+        train_script_path = write_train_script(destination_root, dataset_yaml, DEFAULT_MODEL_SPEC)
+
+    now_utc = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+    class_name_map = {{str(i): name for i, name in enumerate(CLASS_NAMES)}}
+    for class_id in class_counts:
+        class_name_map.setdefault(str(class_id), f"class_{{class_id}}")
+
+    metadata = {{
+        "prepared_at_utc": now_utc,
+        "source_root": str(source_root),
+        "destination_root": str(destination_root),
+        "images_dir_resolved": str(images_dir),
+        "labels_dir_resolved": str(labels_dir),
+        "dry_run": args.dry_run,
+        "train_ratio": args.train_ratio,
+        "seed": args.seed,
+        "dataset_yaml": str(dataset_yaml) if dataset_yaml else None,
+        "train_script": str(train_script_path) if train_script_path else None,
+        "model_spec": DEFAULT_MODEL_SPEC,
+        "class_names": class_name_map,
+        "class_counts": {{str(k): int(v) for k, v in sorted(class_counts.items())}},
+        "totals": {{
+            "pairs": len(pairs),
+            "train_pairs": len(train_pairs),
+            "val_pairs": len(val_pairs),
+            "missing_labels": len(missing_labels),
+            "orphan_labels": len(orphan_labels),
+        }},
+        "train_files": train_rows,
+        "val_files": val_rows,
+        "missing_labels": missing_labels,
+        "orphan_labels": orphan_labels,
+    }}
+
+    logs_root = destination_root / "logs"
+    if not args.dry_run:
+        logs_root.mkdir(parents=True, exist_ok=True)
+        stamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
+        json_log = logs_root / f"dataset_preparation_{{stamp}}.json"
+        txt_log = logs_root / f"dataset_preparation_{{stamp}}.txt"
+        json_log.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+        txt_log.write_text(build_text_log(metadata), encoding="utf-8")
+        print(f"Wrote logs: {{json_log}} and {{txt_log}}")
+        print(f"Wrote training script: {{train_script_path}}")
+    else:
+        print("Dry run enabled: files were not copied and logs were not written.")
+        print("Summary:")
+        print(json.dumps(metadata["totals"], indent=2))
+
+    print(
+        f"Prepared dataset split from {{source_root}} -> {{destination_root}} "
+        f"(train={{len(train_pairs)}}, val={{len(val_pairs)}}, total={{len(pairs)}})."
+    )
+
+
+if __name__ == "__main__":
+    main()
+'''
+        return textwrap.dedent(script)

src/gui/review_mixin.py

@@ -0,0 +1,474 @@
+"""Review dataset workflow mixin (YOLO/COCO)."""
+
+from pathlib import Path
+import json
+
+import cv2
+import numpy as np
+from tkinter import filedialog, messagebox
+
+
+class ReviewMixin:
+    """Methods for opening, navigating, parsing, and saving review datasets."""
+
+    def open_review_folder(self):
+        """Open local YOLO or COCO review dataset."""
+        selected = filedialog.askdirectory(title="Select review root (or images folder)")
+        if not selected:
+            return
+
+        root = Path(selected)
+        coco_candidates = [root / "annotations.json", root / "coco" / "annotations.json"]
+        coco_path = next((p for p in coco_candidates if p.exists()), None)
+
+        if coco_path is not None:
+            try:
+                coco = json.loads(coco_path.read_text(encoding="utf-8"))
+            except Exception as e:
+                messagebox.showerror("COCO Load Error", f"Failed reading COCO file:\n{e}")
+                return
+
+            images = coco.get("images", [])
+            annotations = coco.get("annotations", [])
+            anns_by_image = {}
+            for ann in annotations:
+                anns_by_image.setdefault(ann.get("image_id"), []).append(ann)
+
+            base_dir = coco_path.parent
+            image_search_dirs = [base_dir / "images", root / "images", base_dir, root]
+            items = []
+            for img in images:
+                file_name = str(img.get("file_name", "")).strip()
+                if not file_name:
+                    continue
+                found_path = None
+                for d in image_search_dirs:
+                    candidate = d / file_name
+                    if candidate.exists():
+                        found_path = candidate
+                        break
+                if found_path is None:
+                    continue
+                items.append({
+                    "format": "coco",
+                    "image_id": img.get("id"),
+                    "image_path": found_path,
+                    "file_name": file_name,
+                    "annotations": anns_by_image.get(img.get("id"), []),
+                })
+
+            if not items:
+                messagebox.showwarning(
+                    "No COCO Images",
+                    "Found annotations.json but no readable image files referenced by COCO images[].",
+                )
+                return
+
+            self._stop_streaming()
+            self.review_mode = True
+            self.review_format = "coco"
+            self.review_coco_path = coco_path
+            self.review_coco_data = coco
+            self.review_items = sorted(items, key=lambda it: str(it["image_path"]))
+            self.review_index = 0
+            self.current_review_item = None
+            self.load_review_item(0)
+            return
+
+        images_dir = root / "images" if (root / "images").is_dir() else root
+        labels_dir = root / "labels" if (root / "labels").is_dir() else root
+        image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
+
+        items = []
+        for img_path in sorted(images_dir.rglob("*")):
+            if not img_path.is_file() or img_path.suffix.lower() not in image_exts:
+                continue
+            rel = img_path.relative_to(images_dir)
+            label_path = (labels_dir / rel).with_suffix(".txt")
+            if not label_path.exists():
+                label_path = labels_dir / f"{img_path.stem}.txt"
+            if label_path.exists():
+                items.append({"image_path": img_path, "label_path": label_path})
+
+        if not items:
+            messagebox.showwarning(
+                "No Pairs Found",
+                "No image+label pairs found. Expected either:\n"
+                "1) <root>/images + <root>/labels with matching names\n"
+                "2) images and .txt labels in same folder",
+            )
+            return
+
+        self._stop_streaming()
+        self.review_mode = True
+        self.review_format = "yolo"
+        self.review_coco_path = None
+        self.review_coco_data = None
+        self.review_items = items
+        self.review_index = 0
+        self.current_review_item = None
+        self.load_review_item(0)
+
+    def exit_review_mode(self, silent=False):
+        """Exit local review mode and clear review state."""
+        was_active = self.review_mode
+        self.review_mode = False
+        self.review_format = None
+        self.review_coco_path = None
+        self.review_coco_data = None
+        self.review_items = []
+        self.review_index = -1
+        self.current_review_item = None
+        if was_active and not silent:
+            self.status_var.set("Exited review mode")
+
+    def prev_review_image(self):
+        """Load previous review image."""
+        if not self.review_mode or not self.review_items:
+            messagebox.showwarning("Review Mode", "Open a review folder first.")
+            return
+        if self.review_index <= 0:
+            self.status_var.set("Already at first review image")
+            return
+        self.load_review_item(self.review_index - 1)
+
+    def next_review_image(self):
+        """Load next review image."""
+        if not self.review_mode or not self.review_items:
+            messagebox.showwarning("Review Mode", "Open a review folder first.")
+            return
+        if self.review_index >= len(self.review_items) - 1:
+            self.status_var.set("Already at last review image")
+            return
+        self.load_review_item(self.review_index + 1)
+
+    def _on_review_next_shortcut(self, event=None):
+        """Keyboard shortcut: next image in review mode."""
+        if not self.review_mode:
+            return
+        self.next_review_image()
+
+    def _on_review_prev_shortcut(self, event=None):
+        """Keyboard shortcut: previous image in review mode."""
+        if not self.review_mode:
+            return
+        self.prev_review_image()
+
+    def _parse_yolo_label_to_detections(self, label_path: Path, img_w: int, img_h: int):
+        """Parse YOLO label file (segmentation preferred, detection tolerated) into local detections."""
+        detections = []
+        try:
+            lines = label_path.read_text(encoding="utf-8").splitlines()
+        except Exception:
+            return detections
+
+        for line in lines:
+            stripped = line.strip()
+            if not stripped:
+                continue
+            parts = stripped.split()
+            if len(parts) < 5:
+                continue
+
+            try:
+                class_id = int(float(parts[0]))
+                vals = [float(v) for v in parts[1:]]
+            except Exception:
+                continue
+
+            pts = []
+            if len(vals) >= 6 and len(vals) % 2 == 0:
+                for i in range(0, len(vals), 2):
+                    x = min(max(vals[i] * img_w, 0.0), img_w - 1.0)
+                    y = min(max(vals[i + 1] * img_h, 0.0), img_h - 1.0)
+                    pts.append([x, y])
+            elif len(vals) == 4:
+                # bbox line fallback -> convert to rectangle polygon
+                cx, cy, bw, bh = vals
+                x1 = (cx - bw / 2.0) * img_w
+                y1 = (cy - bh / 2.0) * img_h
+                x2 = (cx + bw / 2.0) * img_w
+                y2 = (cy + bh / 2.0) * img_h
+                x1 = min(max(x1, 0.0), img_w - 1.0)
+                y1 = min(max(y1, 0.0), img_h - 1.0)
+                x2 = min(max(x2, 0.0), img_w - 1.0)
+                y2 = min(max(y2, 0.0), img_h - 1.0)
+                pts = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
+            else:
+                continue
+
+            arr = np.array(pts, dtype=np.float32)
+            x1, y1 = float(np.min(arr[:, 0])), float(np.min(arr[:, 1]))
+            x2, y2 = float(np.max(arr[:, 0])), float(np.max(arr[:, 1]))
+            poly = (arr - [x1, y1]).reshape(-1, 1, 2).astype(np.float32)
+            class_name = self.id_to_class.get(class_id, self._default_class_name())
+
+            detections.append(
+                {
+                    "box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
+                    "name": class_name,
+                    "poly": poly,
+                    "polygons": [poly],
+                    "conf": 1.0,
+                }
+            )
+
+        return detections
+
+    def _parse_coco_item_to_detections(self, item):
+        """Parse COCO image item into local detections with multi-island polygons."""
+        detections = []
+        coco = self.review_coco_data or {}
+        categories = coco.get("categories", [])
+        cat_by_id = {cat.get("id"): cat.get("name", f"class_{cat.get('id')}") for cat in categories}
+        img = cv2.imread(str(item.get("image_path", "")))
+
+        for ann in item.get("annotations", []):
+            seg = ann.get("segmentation", [])
+            polygons_abs = []
+            holes_abs = []
+
+            if isinstance(seg, dict):
+                mask = self._rle_decode_binary_mask(seg)
+                if mask is not None:
+                    contours, hierarchy = cv2.findContours((mask * 255).astype(np.uint8), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
+                    if hierarchy is not None and len(contours) > 0:
+                        for i, cnt in enumerate(contours):
+                            arr = cnt.reshape(-1, 2).astype(np.float32)
+                            if len(arr) < 3:
+                                continue
+                            parent = hierarchy[0][i][3]
+                            if parent == -1:
+                                polygons_abs.append(arr)
+                            else:
+                                holes_abs.append(arr)
+            else:
+                if not isinstance(seg, list):
+                    seg = []
+                for poly in seg:
+                    if not isinstance(poly, list) or len(poly) < 6 or (len(poly) % 2) != 0:
+                        continue
+                    pts = []
+                    for i in range(0, len(poly), 2):
+                        try:
+                            x = float(poly[i])
+                            y = float(poly[i + 1])
+                        except Exception:
+                            pts = []
+                            break
+                        pts.append([x, y])
+                    if len(pts) >= 3:
+                        polygons_abs.append(np.array(pts, dtype=np.float32))
+
+            if not polygons_abs:
+                bbox = ann.get("bbox", [0, 0, 0, 0])
+                if isinstance(bbox, list) and len(bbox) >= 4:
+                    x, y, w, h = [float(v) for v in bbox[:4]]
+                    polygons_abs = [np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]], dtype=np.float32)]
+                else:
+                    continue
+
+            all_pts = np.vstack(polygons_abs)
+            x1 = float(np.min(all_pts[:, 0]))
+            y1 = float(np.min(all_pts[:, 1]))
+            x2 = float(np.max(all_pts[:, 0]))
+            y2 = float(np.max(all_pts[:, 1]))
+            rel_polygons = [
+                (poly - [x1, y1]).reshape(-1, 1, 2).astype(np.float32)
+                for poly in polygons_abs
+            ]
+            rel_holes = [
+                (poly - [x1, y1]).reshape(-1, 1, 2).astype(np.float32)
+                for poly in holes_abs
+            ]
+            class_id = ann.get("category_id")
+            try:
+                class_id_int = int(class_id)
+            except Exception:
+                class_id_int = 1
+            class_name = cat_by_id.get(class_id, self.id_to_class.get(class_id_int - 1, self._default_class_name()))
+
+            detections.append(
+                {
+                    "box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
+                    "name": class_name,
+                    "poly": rel_polygons[0],
+                    "polygons": rel_polygons,
+                    "holes": rel_holes,
+                    "conf": 1.0,
+                    "_review_category_id": class_id,
+                    "_review_ann_id": ann.get("id"),
+                }
+            )
+
+        return detections
+
+    def load_review_item(self, index: int):
+        """Load review image and corresponding local annotation."""
+        if index < 0 or index >= len(self.review_items):
+            return
+
+        item = self.review_items[index]
+        img = cv2.imread(str(item["image_path"]))
+        if img is None:
+            messagebox.showerror("Load Error", f"Failed to read image:\n{item['image_path']}")
+            return
+
+        self.current_raw_frame = img
+        self.current_image_id = None
+        self.current_annotation_id = None
+        self.current_local_image_path = item["image_path"]
+        if item.get("format") == "coco" or self.review_format == "coco":
+            self.last_results = self._parse_coco_item_to_detections(item)
+        else:
+            self.last_results = self._parse_yolo_label_to_detections(
+                item["label_path"], img.shape[1], img.shape[0]
+            )
+        self.canvas.selected_obj_idx = -1
+        self.selected_island_idx = 0
+        self.selected_island_kind = "outer"
+        self.review_index = index
+        self.current_review_item = item
+        self.redraw_current_frame()
+        mode = "COCO" if (item.get("format") == "coco" or self.review_format == "coco") else "YOLO"
+        self.status_var.set(
+            f"Review {self.review_index + 1}/{len(self.review_items)} [{mode}]: {item['image_path'].name}"
+        )
+
+    def _build_segmentation_lines_from_current(self):
+        """Build YOLO segmentation lines from current detections (polygon-only output format)."""
+        if self.current_raw_frame is None:
+            return []
+        img_h, img_w = self.current_raw_frame.shape[:2]
+        lines = []
+        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)
+            polygons = self._det_get_absolute_polygons(det)
+            if not polygons:
+                x1, y1 = float(box.get("x1", 0)), float(box.get("y1", 0))
+                x2, y2 = float(box.get("x2", 0)), float(box.get("y2", 0))
+                polygons = [np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.float32)]
+
+            for points in polygons:
+                norm = []
+                for pt in points:
+                    nx = float(np.clip(pt[0] / img_w, 0.0, 1.0))
+                    ny = float(np.clip(pt[1] / img_h, 0.0, 1.0))
+                    norm.extend([nx, ny])
+                lines.append(f"{class_id} " + " ".join(f"{p:.6f}" for p in norm))
+        return lines
+
+    def save_review_annotation(self):
+        """Save current edited annotation back to local review label file."""
+        if not self.review_mode or not self.current_review_item:
+            messagebox.showwarning("Review Mode", "No active review item.")
+            return
+        if self.current_raw_frame is None:
+            messagebox.showwarning("Warning", "No image loaded.")
+            return
+
+        if self.current_review_item.get("format") == "coco" or self.review_format == "coco":
+            self._save_review_coco_annotation()
+            return
+
+        lines = self._build_segmentation_lines_from_current()
+        label_path = self.current_review_item["label_path"]
+        try:
+            label_path.parent.mkdir(parents=True, exist_ok=True)
+            label_path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
+            self.status_var.set(f"Saved review annotation: {label_path.name}")
+        except Exception as e:
+            messagebox.showerror("Save Error", f"Failed writing label:\n{e}")
+
+    def _save_review_coco_annotation(self):
+        """Save current review image annotations back to COCO annotations.json."""
+        item = self.current_review_item or {}
+        image_id = item.get("image_id")
+        if image_id is None:
+            messagebox.showerror("Save Error", "COCO review item missing image_id.")
+            return
+        if self.review_coco_path is None or self.review_coco_data is None:
+            messagebox.showerror("Save Error", "No active COCO review dataset loaded.")
+            return
+
+        try:
+            coco = self.review_coco_data
+            coco.setdefault("images", [])
+            coco.setdefault("annotations", [])
+            coco.setdefault("categories", [])
+            categories = coco.get("categories", [])
+            cat_by_name = {str(cat.get("name")): cat.get("id") for cat in categories if cat.get("name")}
+            max_cat_id = max((int(cat.get("id", 0)) for cat in categories), default=0)
+
+            coco["annotations"] = [ann for ann in coco["annotations"] if ann.get("image_id") != image_id]
+            next_ann_id = max((int(ann.get("id", 0)) for ann in coco["annotations"]), default=0) + 1
+
+            for det in self.last_results:
+                class_name = str(det.get("name", det.get("label", self._default_class_name())))
+                category_id = cat_by_name.get(class_name)
+                if category_id is None:
+                    stored_cat_id = det.get("_review_category_id")
+                    if isinstance(stored_cat_id, int):
+                        category_id = stored_cat_id
+                    else:
+                        max_cat_id += 1
+                        category_id = max_cat_id
+                        categories.append({"id": category_id, "name": class_name})
+                    cat_by_name[class_name] = category_id
+
+                box = det.get("box", det)
+                x1, y1 = float(box.get("x1", 0.0)), float(box.get("y1", 0.0))
+                x2, y2 = float(box.get("x2", 0.0)), float(box.get("y2", 0.0))
+                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:
+                    img_h, img_w = self.current_raw_frame.shape[:2]
+                    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": int(category_id),
+                        "segmentation": segmentation_payload,
+                        "area": float(area_total),
+                        "bbox": bbox_payload,
+                        "iscrowd": 0,
+                    }
+                )
+                next_ann_id += 1
+
+            # Keep in-memory review item synchronized for immediate reload/navigation.
+            image_anns = [ann for ann in coco["annotations"] if ann.get("image_id") == image_id]
+            self.current_review_item["annotations"] = image_anns
+            self.review_coco_path.write_text(json.dumps(coco, indent=2), encoding="utf-8")
+            self.status_var.set(f"Saved COCO review annotation: {self.current_review_item['image_path'].name}")
+        except Exception as e:
+            messagebox.showerror("Save Error", f"Failed writing COCO annotations:\n{e}")

src/inference/z_m_q_inference_monitor.py

@@ -0,0 +1,61 @@
+import zmq
+import json
+import time
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(description="Monitor YOLO ZMQ Inference Stream")
+    parser.add_argument('--addr', default='tcp://localhost:9091', help='ZMQ PUB address to subscribe to')
+    parser.add_argument('--topic', default='detections', help='Topic to subscribe to (empty for all)')
+    args = parser.parse_args()
+
+    ctx = zmq.Context()
+    sub = ctx.socket(zmq.SUB)
+    
+    # Configure for low latency
+    sub.setsockopt(zmq.RCVHWM, 1000)
+    sub.connect(args.addr)
+    sub.setsockopt_string(zmq.SUBSCRIBE, args.topic)
+
+    print(f"📡 Monitoring detections on {args.addr} (Topic: '{args.topic}')...")
+
+    msg_count = 0
+    start_time = time.time()
+
+    try:
+        while True:
+            # Check for multipart (topic + json) or single part (just json)
+            parts = sub.recv_multipart()
+            msg_count += 1
+        
+            # FIX: Logic to find the JSON header in a multipart message
+            # Message might be [Topic, JSON, Pixels] or just [JSON, Pixels]
+            header_data = None
+            for p in parts[:-1]: # Don't look at the last part (pixels)
+                try:
+                    decoded = p.decode('utf-8')
+                    if '{' in decoded: # Likely JSON
+                        header_data = json.loads(decoded)
+                        break
+                except: continue
+
+            if header_data:
+                boxes = header_data.get('boxes', [])
+                timestamp = header_data.get('time', 0)
+                latency = (time.time() - timestamp) * 1000 if timestamp > 0 else 0
+
+                print(f"[{msg_count:04d}] Latency: {latency:6.1f}ms | Detections: {len(boxes)} | Img: {header_data.get('shape')}")
+                for det in boxes:
+                    print(f"  └─ Class: {det['label']} Conf: {det['conf']:.2f}")
+            else:
+                print(f"[{msg_count:04d}] Error: Could not find valid JSON header in {len(parts)} parts")
+
+    except KeyboardInterrupt:
+        elapsed = time.time() - start_time
+        print(f"\nStopped. Processed {msg_count} messages in {elapsed:.1f}s ({msg_count/elapsed:.1f} msg/s)")
+    finally:
+        sub.close()
+        ctx.term()
+
+if __name__ == "__main__":
+    main()

src/server/focus_epics.py

@@ -0,0 +1,196 @@
+import queue
+import time
+from threading import Event, Lock, Thread
+from typing import Any
+
+import cv2
+import numpy as np
+
+
+def compute_focus_score_bgr(img_bgr: np.ndarray, scale: float = 0.5) -> float:
+    """
+    Fast focus proxy using variance of Laplacian.
+    Lower scale reduces CPU cost substantially with modest sensitivity tradeoff.
+    """
+    if img_bgr is None or img_bgr.size == 0:
+        return 0.0
+    work = img_bgr
+    s = float(scale)
+    if 0.0 < s < 1.0:
+        h, w = img_bgr.shape[:2]
+        nw = max(1, int(w * s))
+        nh = max(1, int(h * s))
+        work = cv2.resize(work, (nw, nh), interpolation=cv2.INTER_AREA)
+    gray = cv2.cvtColor(work, cv2.COLOR_BGR2GRAY)
+    return float(cv2.Laplacian(gray, cv2.CV_64F).var())
+
+
+class FocusPvWriter:
+    """Non-blocking EPICS writer for focus scores."""
+
+    def __init__(self, pv_name: str, enabled: bool, min_period_ms: float = 100.0):
+        self._lock = Lock()
+        self._queue = queue.Queue(maxsize=1)
+        self._stop = Event()
+        self._thread = Thread(target=self._worker, daemon=True)
+        self._enabled = bool(enabled)
+        self._pv_name = str(pv_name or "").strip()
+        self._min_period_s = max(0.0, float(min_period_ms) / 1000.0)
+        self._last_write_ts = 0.0
+        self._last_write_value = None
+        self._last_error = ""
+        self._submit_count = 0
+        self._drop_count = 0
+        self._write_count = 0
+        self._last_log_ts = 0.0
+        self._backend = "none"
+        self._epics = None
+        self._pv = None
+        self._caproto_write = None
+        self._thread.start()
+
+    def _log(self, msg: str, min_interval_s: float = 5.0) -> None:
+        now = time.time()
+        if (now - self._last_log_ts) >= min_interval_s:
+            print(f"[FocusPvWriter] {msg}", flush=True)
+            self._last_log_ts = now
+
+    def _connect_if_needed(self) -> bool:
+        if self._backend == "caproto" and self._caproto_write is not None:
+            return True
+        if self._pv is not None:
+            return True
+        if not self._pv_name:
+            self._last_error = "focus_pv is empty"
+            self._log("focus_pv is empty; EPICS writes disabled", min_interval_s=10.0)
+            return False
+        try:
+            if self._epics is None:
+                import epics  # type: ignore
+                self._epics = epics
+            self._pv = self._epics.PV(self._pv_name, auto_monitor=False)
+            self._backend = "pyepics"
+            self._log(f"connected to PV '{self._pv_name}'", min_interval_s=0.0)
+            return True
+        except Exception as e:
+            self._last_error = str(e)
+            self._pv = None
+            self._log(
+                f"pyepics failed for '{self._pv_name}': {e}. Trying caproto fallback.",
+                min_interval_s=5.0,
+            )
+            try:
+                from caproto.sync.client import write as caproto_write  # type: ignore
+
+                self._caproto_write = caproto_write
+                self._backend = "caproto"
+                self._log(
+                    f"using caproto fallback backend for PV '{self._pv_name}'",
+                    min_interval_s=0.0,
+                )
+                return True
+            except Exception as caproto_exc:
+                self._last_error = f"pyepics: {e}; caproto: {caproto_exc}"
+                self._backend = "none"
+                self._log(
+                    f"caproto fallback failed for '{self._pv_name}': {caproto_exc}",
+                    min_interval_s=5.0,
+                )
+                return False
+
+    def configure(
+        self,
+        *,
+        enabled: bool | None = None,
+        pv_name: str | None = None,
+        min_period_ms: float | None = None,
+    ) -> None:
+        with self._lock:
+            if enabled is not None:
+                self._enabled = bool(enabled)
+            if pv_name is not None:
+                new_name = str(pv_name or "").strip()
+                if new_name != self._pv_name:
+                    self._pv_name = new_name
+                    self._pv = None
+                    self._caproto_write = None
+                    self._backend = "none"
+            if min_period_ms is not None:
+                self._min_period_s = max(0.0, float(min_period_ms) / 1000.0)
+
+    def submit(self, value: float) -> None:
+        with self._lock:
+            enabled = self._enabled
+        if not enabled:
+            return
+        self._submit_count += 1
+        try:
+            self._queue.put_nowait(float(value))
+        except queue.Full:
+            self._drop_count += 1
+            try:
+                _ = self._queue.get_nowait()
+            except Exception:
+                pass
+            try:
+                self._queue.put_nowait(float(value))
+            except Exception:
+                pass
+
+    def status(self) -> dict[str, Any]:
+        with self._lock:
+            enabled = self._enabled
+            pv_name = self._pv_name
+            min_period_ms = self._min_period_s * 1000.0
+        return {
+            "enabled": bool(enabled),
+            "pv_name": pv_name,
+            "min_period_ms": float(min_period_ms),
+            "backend": self._backend,
+            "connected": bool(self._pv is not None) or bool(self._backend == "caproto"),
+            "submit_count": int(self._submit_count),
+            "drop_count": int(self._drop_count),
+            "write_count": int(self._write_count),
+            "last_write_ts": float(self._last_write_ts),
+            "last_write_value": self._last_write_value,
+            "last_error": self._last_error,
+        }
+
+    def _worker(self) -> None:
+        while not self._stop.is_set():
+            try:
+                val = float(self._queue.get(timeout=0.2))
+            except queue.Empty:
+                continue
+            with self._lock:
+                enabled = self._enabled
+                min_period_s = self._min_period_s
+            if not enabled:
+                continue
+            now = time.time()
+            if min_period_s > 0.0 and (now - self._last_write_ts) < min_period_s:
+                continue
+            if not self._connect_if_needed():
+                continue
+            try:
+                if self._backend == "caproto" and self._caproto_write is not None:
+                    self._caproto_write(self._pv_name, val, notify=False, timeout=0.5)
+                else:
+                    self._pv.put(val, wait=False)
+                self._last_write_ts = now
+                self._last_write_value = val
+                self._write_count += 1
+            except Exception as e:
+                self._last_error = str(e)
+                self._pv = None
+                if self._backend == "caproto":
+                    self._caproto_write = None
+                self._backend = "none"
+                self._log(f"write failed for '{self._pv_name}': {e}", min_interval_s=5.0)
+
+    def close(self) -> None:
+        self._stop.set()
+        try:
+            self._thread.join(timeout=1.0)
+        except Exception:
+            pass

src/server/image_utils.py

@@ -0,0 +1,89 @@
+import cv2
+import numpy as np
+
+
+# per-class color mapping (BGR tuples for OpenCV)
+# keep colors consistent with your labeling tool:
+# class 0 = GREEN, 1 = RED, 2 = BLUE, 3 = YELLOW
+CLASS_COLORS = {
+    0: (0, 255, 0),   # green
+    1: (0, 0, 255),   # red
+    2: (255, 0, 0),   # blue
+    3: (0, 255, 255), # yellow (cyan in BGR)
+}
+
+
+def decode_image(msg_bytes):
+    # try to decode raw JPEG/PNG bytes robustly
+    try:
+        # handle memoryview/bytes/bytearray
+        if isinstance(msg_bytes, memoryview):
+            buf = msg_bytes.tobytes()
+        else:
+            buf = bytes(msg_bytes)
+        arr = np.frombuffer(buf, dtype=np.uint8)
+        # ensure contiguous writable array for OpenCV
+        arr = arr.copy()
+        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+        if img is not None:
+            return img
+    except Exception:
+        pass
+
+    # fallback: if the payload is base64-encoded text, try to decode
+    try:
+        import base64
+
+        b = base64.b64decode(msg_bytes)
+        arr = np.frombuffer(b, dtype=np.uint8).copy()
+        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+        return img
+    except Exception:
+        return None
+
+
+def draw_boxes(img, results, conf_threshold=0.25):
+    # results: ultralytics.Results for single image
+    if results is None:
+        return img
+    boxes = getattr(results, 'boxes', None)
+    if boxes is None or len(boxes) == 0:
+        return img
+
+    names = results.names if hasattr(results, 'names') else {}
+    try:
+        xyxy = boxes.xyxy.cpu().numpy()
+        confs = boxes.conf.cpu().numpy()
+        classes = boxes.cls.cpu().numpy().astype(int)
+    except Exception:
+        # fallback to converting to numpy via .numpy()
+        arr = results.boxes.numpy()
+        # arr columns: x1,y1,x2,y2,score,class
+        xyxy = arr[:, :4]
+        confs = arr[:, 4]
+        classes = arr[:, 5].astype(int)
+
+    for (x1, y1, x2, y2), conf, cls in zip(xyxy, confs, classes):
+        if conf < conf_threshold:
+            continue
+        x1i, y1i, x2i, y2i = int(x1), int(y1), int(x2), int(y2)
+        label = names.get(cls, str(cls))
+        # pick color from mapping, fallback to a deterministic hashed color
+        try:
+            color = CLASS_COLORS.get(int(cls))
+            if color is None:
+                raise KeyError()
+        except Exception:
+            # deterministic fallback: create color from class id
+            v = int(cls)
+            # simple hashing to BGR range
+            b = (37 * v) % 256
+            g = (17 * v) % 256
+            r = (97 * v) % 256
+            color = (b, g, r)
+        cv2.rectangle(img, (x1i, y1i), (x2i, y2i), color, 2)
+        text = f"{label} {conf:.2f}"
+        ((tw, _th), _) = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        cv2.rectangle(img, (x1i, y1i - 20), (x1i + tw, y1i), color, -1)
+        cv2.putText(img, text, (x1i, y1i - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
+    return img

src/server/model_utils.py

@@ -0,0 +1,128 @@
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from ultralytics import YOLO
+
+
+def resolve_task(task: str, pt_path: str, engine_path: str) -> str:
+    """Resolve model task from explicit arg or filename hint."""
+    if task in ("detect", "segment"):
+        return task
+    probe = f"{engine_path or ''} {pt_path or ''}".lower()
+    if "seg" in probe:
+        return "segment"
+    return "detect"
+
+
+def project_root() -> Path:
+    """
+    Resolve repository root robustly even if this module is moved.
+    Prefer the closest ancestor containing a 'models' directory.
+    """
+    here = Path(__file__).resolve()
+    for candidate in here.parents:
+        if (candidate / "models").exists():
+            return candidate
+    # Fallback keeps previous behavior if the models directory does not exist yet.
+    return here.parent.parent
+
+
+def resolve_model_path(path_value: str) -> tuple[str, list[str]]:
+    """
+    Resolve a model path robustly across differing working directories.
+    Returns: (resolved_path_or_raw, tried_paths)
+    """
+    raw = (path_value or "").strip()
+    if not raw:
+        return "", []
+
+    p = Path(raw).expanduser()
+    tried: list[Path] = [p]
+
+    if p.exists():
+        return str(p), [str(x) for x in tried]
+
+    root = project_root()
+    models_dir = root / "models"
+
+    root_candidate = (root / p)
+    tried.append(root_candidate)
+    if root_candidate.exists():
+        return str(root_candidate.resolve()), [str(x) for x in tried]
+
+    raw_norm = raw.replace("\\", "/")
+    if raw_norm.startswith("models/"):
+        rel_in_models = raw_norm[len("models/"):]
+        model_candidate = models_dir / rel_in_models
+        tried.append(model_candidate)
+        if model_candidate.exists():
+            return str(model_candidate.resolve()), [str(x) for x in tried]
+
+    basename = Path(raw_norm).name
+    if basename and models_dir.exists():
+        matches = sorted(x for x in models_dir.rglob(basename) if x.is_file())
+        for m in matches:
+            tried.append(m)
+        if len(matches) == 1:
+            return str(matches[0].resolve()), [str(x) for x in tried]
+
+    return raw, [str(x) for x in tried]
+
+
+def is_segmentation_capable_name(name: str) -> bool:
+    lowered = (name or "").lower()
+    return any(token in lowered for token in ("seg", "segment", "mask"))
+
+
+def list_available_models() -> list[dict[str, Any]]:
+    root = project_root()
+    models_dir = root / "models"
+    if not models_dir.exists():
+        return []
+
+    valid_suffixes = {".pt", ".engine", ".engines", ".yaml", ".yml", ".onnx"}
+    items: list[dict[str, Any]] = []
+    for p in sorted(models_dir.rglob("*")):
+        if not p.is_file() or p.suffix.lower() not in valid_suffixes:
+            continue
+        rel_from_models = str(p.relative_to(models_dir))
+        rel_from_root = str(p.relative_to(root))
+        segmentation_capable = is_segmentation_capable_name(p.name)
+        items.append(
+            {
+                "name": rel_from_models,
+                "path": str(p),
+                "project_path": rel_from_root,
+                "kind": p.suffix.lower().lstrip("."),
+                "segmentation_capable": segmentation_capable,
+                "task_hint": "segment" if segmentation_capable else "detect",
+            }
+        )
+    return items
+
+
+def load_model(pt_path: str, engine_path: str, task: str):
+    resolved_pt, pt_tried = resolve_model_path(pt_path)
+    resolved_engine, engine_tried = resolve_model_path(engine_path)
+    pt_path = resolved_pt
+    engine_path = resolved_engine
+
+    resolved_task = resolve_task(task, pt_path, engine_path)
+    if engine_path and os.path.exists(engine_path):
+        print(f"Loading TensorRT engine: {engine_path}")
+        print(f"Model task: {resolved_task}")
+        return YOLO(engine_path, task=resolved_task)
+    if pt_path and os.path.exists(pt_path):
+        print(f"Loading PyTorch model: {pt_path}")
+        print(f"Model task: {resolved_task}")
+        return YOLO(pt_path, task=resolved_task)
+    debug = {
+        "engine": engine_path,
+        "pt": pt_path,
+        "engine_tried": engine_tried,
+        "pt_tried": pt_tried,
+        "project_root": str(project_root()),
+    }
+    raise FileNotFoundError(f"No valid model found. Resolution details: {json.dumps(debug)}")

src/server/runtime_config.py

@@ -0,0 +1,42 @@
+import json
+from pathlib import Path
+from typing import Any
+
+
+def load_runtime_config(path: str) -> tuple[dict[str, Any], str | None]:
+    p = Path(path).expanduser()
+    if not p.exists():
+        return {}, None
+    try:
+        data = json.loads(p.read_text(encoding="utf-8"))
+    except Exception as e:
+        return {}, f"Failed to parse config '{p}': {e}"
+    if not isinstance(data, dict):
+        return {}, f"Config '{p}' must be a JSON object"
+    return data, None
+
+
+def save_runtime_config(path: str, config: dict[str, Any]) -> str | None:
+    p = Path(path).expanduser()
+    try:
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_text(json.dumps(config, indent=2, sort_keys=True), encoding="utf-8")
+        return None
+    except Exception as e:
+        return f"Failed to save config '{p}': {e}"
+
+
+def collect_cli_overrides(argv: list[str]) -> set[str]:
+    """
+    Collect option names explicitly passed on CLI, normalized to argparse dest form.
+    Example: --focus-epics-enabled -> focus_epics_enabled
+    """
+    out: set[str] = set()
+    for token in argv[1:]:
+        if not token.startswith("--"):
+            continue
+        key = token[2:]
+        if "=" in key:
+            key = key.split("=", 1)[0]
+        out.add(key.replace("-", "_"))
+    return out

src/server/stream_utils.py

@@ -0,0 +1,70 @@
+import queue
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from threading import Lock, Thread
+
+import cv2
+
+
+_latest_frame = None
+_frame_lock = Lock()
+_encode_queue = queue.Queue(maxsize=2)
+
+
+class MJPEGHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path != '/stream' and self.path != '/':
+            self.send_response(200)
+            self.send_header('Content-type', 'text/html')
+            self.end_headers()
+            self.wfile.write(b"<html><body><img src='/stream' /></body></html>")
+            return
+        self.send_response(200)
+        self.send_header('Content-type', 'multipart/x-mixed-replace; boundary=FRAME')
+        self.end_headers()
+        try:
+            while True:
+                with _frame_lock:
+                    frame = _latest_frame
+                if frame:
+                    self.wfile.write(b'--FRAME\r\n')
+                    self.wfile.write(b'Content-Type: image/jpeg\r\n')
+                    self.wfile.write(f'Content-Length: {len(frame)}\r\n\r\n'.encode())
+                    self.wfile.write(frame)
+                    self.wfile.write(b'\r\n')
+                    try:
+                        self.wfile.flush()
+                    except Exception:
+                        pass
+        except Exception:
+            return
+
+
+def _encoding_worker(quality: int, scale_factor: float):
+    """Background thread to handle JPEG compression without blocking the GPU loop."""
+    global _latest_frame
+    while True:
+        try:
+            frame_to_encode = _encode_queue.get(timeout=1)
+            if scale_factor < 1.0:
+                h, w = frame_to_encode.shape[:2]
+                frame_to_encode = cv2.resize(frame_to_encode, (int(w * scale_factor), int(h * scale_factor)))
+            _, enc = cv2.imencode('.jpg', frame_to_encode, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
+            with _frame_lock:
+                _latest_frame = enc.tobytes()
+        except queue.Empty:
+            continue
+
+
+def start_http_server(port: int, quality: int, scale_factor: float):
+    server = HTTPServer(('', port), MJPEGHandler)
+    Thread(target=server.serve_forever, daemon=True).start()
+    Thread(target=_encoding_worker, args=(quality, scale_factor), daemon=True).start()
+    return server
+
+
+def submit_stream_frame(frame) -> None:
+    """Non-blocking frame submission for HTTP MJPEG stream."""
+    try:
+        _encode_queue.put_nowait(frame)
+    except queue.Full:
+        pass

src/server/target_point.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from typing import Any
+
+import cv2
+import numpy as np
+
+
+def _poly_to_abs(poly: Any, x1: float, y1: float) -> np.ndarray | None:
+    """Convert det['poly'] (relative-to-box) into absolute image points."""
+    if poly is None:
+        return None
+    try:
+        pts = np.asarray(poly, dtype=np.float32)
+    except Exception:
+        return None
+    if pts.ndim != 2 or pts.shape[1] != 2 or pts.shape[0] < 3:
+        return None
+    pts_abs = np.empty_like(pts)
+    pts_abs[:, 0] = pts[:, 0] + float(x1)
+    pts_abs[:, 1] = pts[:, 1] + float(y1)
+    return pts_abs
+
+
+def _shape_center(det: dict[str, Any]) -> tuple[float, float] | None:
+    """Center of mass of polygon shape if present, otherwise bbox center."""
+    try:
+        x1 = float(det.get("x1", 0.0))
+        y1 = float(det.get("y1", 0.0))
+        x2 = float(det.get("x2", 0.0))
+        y2 = float(det.get("y2", 0.0))
+    except Exception:
+        return None
+
+    pts_abs = _poly_to_abs(det.get("poly"), x1, y1)
+    if pts_abs is not None:
+        try:
+            cnt = pts_abs.reshape((-1, 1, 2))
+            m = cv2.moments(cnt)
+            if abs(m["m00"]) > 1e-6:
+                return float(m["m10"] / m["m00"]), float(m["m01"] / m["m00"])
+        except Exception:
+            pass
+
+    return 0.5 * (x1 + x2), 0.5 * (y1 + y2)
+
+
+def _loop_left_area_center(det: dict[str, Any]) -> tuple[float, float] | None:
+    """
+    Target point for loop_all:
+    center of the left-side loop area instead of the top-most apex.
+    """
+    try:
+        x1 = float(det.get("x1", 0.0))
+        y1 = float(det.get("y1", 0.0))
+        x2 = float(det.get("x2", 0.0))
+        y2 = float(det.get("y2", 0.0))
+    except Exception:
+        return None
+
+    pts_abs = _poly_to_abs(det.get("poly"), x1, y1)
+    if pts_abs is not None and len(pts_abs) > 0:
+        min_x = float(np.min(pts_abs[:, 0]))
+        max_x = float(np.max(pts_abs[:, 0]))
+        min_y = float(np.min(pts_abs[:, 1]))
+        max_y = float(np.max(pts_abs[:, 1]))
+
+        width = max_x - min_x
+        height = max_y - min_y
+        if width <= 1e-6 or height <= 1e-6:
+            return 0.5 * (x1 + x2), 0.5 * (y1 + y2)
+
+        y_mid = 0.5 * (min_y + max_y)
+
+        # Keep only points near the left edge of the shape.
+        # This tends to land inside the left "loop area".
+        left_band = pts_abs[pts_abs[:, 0] <= (min_x + 0.35 * width)]
+        if len(left_band) == 0:
+            left_band = pts_abs
+
+        # Prefer points closest to the vertical middle of that left band.
+        idx = int(np.argmin(np.abs(left_band[:, 1] - y_mid)))
+        candidate = left_band[idx]
+
+        # Blend with the left-band centroid for stability.
+        centroid = np.mean(left_band, axis=0)
+        return float(0.5 * (candidate[0] + centroid[0])), float(0.5 * (candidate[1] + centroid[1]))
+
+    return 0.5 * (x1 + x2), 0.5 * (y1 + y2)
+
+
+def _pin_left_midpoint(det: dict[str, Any]) -> tuple[float, float] | None:
+    """Left-middle point of the pin bounding box."""
+    try:
+        x1 = float(det.get("x1", 0.0))
+        y1 = float(det.get("y1", 0.0))
+        y2 = float(det.get("y2", 0.0))
+    except Exception:
+        return None
+    return x1, 0.5 * (y1 + y2)
+
+
+def _pick_best(dets: list[dict[str, Any]], label: str) -> dict[str, Any] | None:
+    class_hits = [d for d in dets if str(d.get("label", "")) == label]
+    if not class_hits:
+        return None
+    return max(class_hits, key=lambda d: float(d.get("conf", 0.0)))
+
+
+def compute_target_point(dets: list[dict[str, Any]]) -> dict[str, Any] | None:
+    """
+    Target priority:
+    1) crystal   -> center of mass of crystal shape
+    2) loop_face -> center of mass of loop_face shape
+    3) loop_all  -> center of left-side loop area
+    4) pin       -> left middle side of pin box
+    """
+    if not dets:
+        return None
+
+    crystal = _pick_best(dets, "crystal")
+    if crystal is not None:
+        pt = _shape_center(crystal)
+        if pt is not None:
+            return {"x": int(round(pt[0])), "y": int(round(pt[1])), "source": "crystal_center"}
+
+    loop_face = _pick_best(dets, "loop_face")
+    if loop_face is not None:
+        pt = _shape_center(loop_face)
+        if pt is not None:
+            return {"x": int(round(pt[0])), "y": int(round(pt[1])), "source": "loop_face_center"}
+
+    loop_all = _pick_best(dets, "loop_all")
+    if loop_all is not None:
+        pt = _loop_left_area_center(loop_all)
+        if pt is not None:
+            return {"x": int(round(pt[0])), "y": int(round(pt[1])), "source": "loop_all_left_area_center"}
+
+    pin = _pick_best(dets, "pin")
+    if pin is not None:
+        pt = _pin_left_midpoint(pin)
+        if pt is not None:
+            return {"x": int(round(pt[0])), "y": int(round(pt[1])), "source": "pin_left_mid"}
+
+    return None

src/tools/alc_integration_yolo.py

@@ -0,0 +1,42 @@
+from ultralytics import YOLO
+import cv2
+
+def get_feature_0_bounding_boxes(image, model_path):
+    """
+    Processes a given image through the YOLO model and extracts the bounding boxes for feature/class 0.
+
+    Args:
+        image: The image object (e.g., an array or PIL.Image supported by YOLO).
+        model_path: Path to the YOLO model file (e.g., best.pt).
+
+    Returns:
+        List of bounding boxes for feature/class 0. Each bounding box is in the format [x1, y1, x2, y2].
+    """
+    # Load the YOLO model
+    model = YOLO(model_path)
+    # Perform prediction on the image
+    results = model.predict(source=image, conf=0.5, save=True, show=False)
+
+    # Extract the first bounding box for class 0
+    for r in results:
+        for box in r.boxes.data.tolist():  # Each box: [x1, y1, x2, y2, confidence, class_id]
+            x1, y1, x2, y2, conf, cls = box
+            if int(cls) == 0:  # Check if class_id is 0 (feature 0)
+                return [x1, y1, x2, y2]  # Return the first box for class 0
+    return None  # If no box is found, return None
+
+
+# Example usage
+if __name__ == "__main__":
+
+    # Provide your YOLO model path and an image
+    yolo_model_path = "/Users/duan_j/Applications/alc/tests/yolo/runs/aare-test/train8/weights/best.pt"  # Replace with your YOLO model's actual path
+    image_path = "/Users/duan_j/Applications/alc/tests/yolo/Snipaste_2025-02-25_14-43-11.jpg"  # Replace with your image's path
+
+    # Open image with PIL (YOLO supports this format)
+    image = cv2.imread(image_path)
+
+    # Get bounding boxes for class 0
+    feature_0_boxes = get_feature_0_bounding_boxes(image, yolo_model_path) #x1, y1, x2, y2
+
+    print("Bounding Boxes for Feature 0:", feature_0_boxes)

src/tools/annotating_imgs.py

@@ -0,0 +1,188 @@
+import cv2
+import numpy as np
+import os
+import glob
+import yaml
+from pathlib import Path
+
+def resize_to_screen(img, screen_width=1920, screen_height=1080):
+    """Resize the image to fit the screen while maintaining the aspect ratio."""
+    height, width, _ = img.shape
+    scaling_factor = min(screen_width / width, screen_height / height)
+    new_width = int(width * scaling_factor)
+    new_height = int(height * scaling_factor)
+    resized_img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    return resized_img, scaling_factor
+
+def draw_and_record_boxes(image_path, screen_width=1920, screen_height=1080):
+    """Draw four boxes and record normalized YOLO box data."""
+    img = cv2.imread(image_path)
+
+    # Resize image to fit the screen
+    resized_img, scale = resize_to_screen(img, screen_width/2, screen_height/2)
+    height, width, _ = img.shape  # Original dimensions (for normalization)
+
+    while True:
+
+        boxes_data = []  # List to store box info
+        preview_img = resized_img.copy()
+
+        # Helper for selecting and recording a box
+        def select_and_record_box(box_name, color, class_id):
+            box = cv2.selectROI(box_name, resized_img, fromCenter=False, showCrosshair=True)
+            if box == (0, 0, 0, 0):  # no selection
+                return None
+            x, y, w, h = [int(coord / scale) for coord in box]  # Scale back to original size
+            rx, ry, rw, rh = [int(coord) for coord in box]  # use resized coords directly
+            cv2.rectangle(preview_img, (rx, ry), (rx + rw, ry + rh), color, 2)
+            #cv2.rectangle(resized_img, (x, y), (x + w, y + h), color, 2)
+            center_x = (x + w / 2) / width
+            center_y = (y + h / 2) / height
+            norm_w = w / width
+            norm_h = h / height
+            return [class_id, center_x, center_y, norm_w, norm_h]
+
+            # Boxes
+
+        box = select_and_record_box(f"Select Class 0: {class_dict[0].upper()} GREEN Box", (0, 255, 0), 0)
+        if box: boxes_data.append(box)
+
+        box = select_and_record_box(f"Select Class 1: {class_dict[1].upper()} RED Box", (0, 0, 255), 1)
+        if box: boxes_data.append(box)
+
+        box = select_and_record_box(f"Select Class 2: {class_dict[2].upper()} BLUE Box", (255, 0, 0), 2)
+        if box: boxes_data.append(box)
+
+        box = select_and_record_box(f"Select Class 3: {class_dict[3].upper()} YELLOW Box", (0, 255, 255), 3)
+        if box: boxes_data.append(box)
+
+        # Show annotated image
+        cv2.imshow("Annotated Image (Press: [a]=accept, [r]=restart, [q/ESC]=quit)", preview_img)
+        key = cv2.waitKey(0) & 0xFF  # Mask to 8-bit
+        cv2.destroyAllWindows()
+
+        if key in [ord('a'), 13]:  # 'a' or ENTER = accept
+            return boxes_data, False, preview_img
+        elif key == ord('r'):  # restart selection
+            print("Restarting annotation for this image...")
+            continue
+        elif key in [27, ord('q')]:  # ESC or 'q' = quit all
+            return boxes_data, True, preview_img
+        elif key == ord('s'):  # s = skip current image
+            print("Skipping this image...")
+            return [], False, preview_img
+        else:
+            print("Unrecognized key. Press [a]=accept, [r]=restart, [q]=quit.")
+
+        return boxes_data, False, preview_img
+
+
+def process_images_for_yolo(input_folder, labels_folder, screen_width=1920, screen_height=1080, overwrite_existing=None):
+    """Process images for YOLO by saving labeled bounding boxes for each image."""
+    # Supported image extensions
+    image_extensions = ('*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tif', '*.tiff')
+
+    # Collect all image file paths from the input folder
+    image_paths = []
+    for ext in image_extensions:
+        image_paths.extend(glob.glob(os.path.join(input_folder, ext)))
+
+    if not image_paths:
+        print(f"No images found in the input folder: {input_folder}")
+        return
+
+    # Create the labels folder if it doesn't exist
+    if not os.path.exists(labels_folder):
+        os.makedirs(labels_folder)
+
+    print("\n--- YOLO Image Annotation Controls ---")
+    print("[ESC] or [q]  → quit annotation completely")
+    print("[c] - cancel process")
+    print("[r]            → restart current image")
+    print("[s]           → skip current image")
+    print("[any other]   → save boxes & move to next image")
+    print("---------------------------------------\n")
+
+    # Iterate over each image
+    for image_path in image_paths:
+        print(f"Processing: {image_path}")
+        image_name = os.path.splitext(os.path.basename(image_path))[0]
+        output_file_path = os.path.join(labels_folder, f"{image_name}.txt")
+
+        if os.path.exists(output_file_path):
+            if overwrite_existing is True:
+                print(f"overwriteing exisiting label for {image_name}.txt")
+            elif overwrite_existing is False:
+                print(f"Skipping {image_name} as already labeled")
+                continue
+            else:
+                print(f"⚠️ Label already exists for {image_name}.txt")
+                choice = input("[o] → overwrite, [s] → skip, [q/ESC] → quit: "). strip().lower()
+
+                if choice == "s":
+                    print(f"Skipping {image_name} (already labeled).")
+                    continue
+                elif choice == "q":
+                    print("Exiting early...")
+                    return
+                elif choice == "o":
+                    print(f"Overwriting label for {image_name}...")
+                else:
+                    print("Unknown choice, skipping this image.")
+                    continue
+
+        print(f"\nProcessing: {image_path}")
+
+        while True:
+            # Get normalized box data from user interaction
+            boxes_data, exit_flag, preview_img = draw_and_record_boxes(image_path, screen_width, screen_height)
+
+            if exit_flag:  # Exit the annotation process if Ctrl+W is pressed
+                print("Exiting image annotation early...")
+                return
+
+            if boxes_data:  # Write the annotations for this image file if boxes are drawn
+                image_name = os.path.splitext(os.path.basename(image_path))[0]  # File name without extension
+                output_file_path = os.path.join(labels_folder, f"{image_name}.txt")
+
+                # Write each box to its corresponding text file in YOLO format
+                with open(output_file_path, 'w') as f:
+                    for box in boxes_data:
+                        class_id, center_x, center_y, norm_w, norm_h = box
+                        f.write(f"{class_id} {center_x:.6f} {center_y:.6f} {norm_w:.6f} {norm_h:.6f}\n")
+
+                #cv2.imshow(f"Saved Annotation: {image_name}", preview_img)
+                #cv2.waitKey(1500)  # wait 1.5 seconds or until key is pressed
+                #cv2.destroyAllWindows()
+
+                break
+
+            else:
+                print("skipping this image...")
+                break
+
+    print(f"Annotations saved to {labels_folder}")
+
+
+
+if __name__ == "__main__":
+    # Define input folder containing images
+    input_folder = r"/Users/gotthardg/Volumes/ra_work/Sept/annotations" #"/Users/duan_j/Applications/alc/tests/yolo/nodetec/pin2"  # Change to your folder path
+    #overwrite flag to say whether to overwrite existing labels, skip all labeled or enable user choice.
+    overwrite_flag = True
+    # Define the labels output folder for YOLO format
+    labels_folder = f"{input_folder}labels"  # Change to your desired folder
+
+    # Path to your yaml file
+    yaml_file = str(Path(__file__).resolve().parents[2] / "config" / "dataset.yaml")
+
+    with open(yaml_file, "r") as f:
+        data = yaml.safe_load(f)
+
+    # Extract names dictionary
+    class_dict = data.get("names", {})
+
+    print(class_dict[0])
+
+    # Process images and save YOLO-compatible annotations
+    process_images_for_yolo(input_folder, labels_folder, overwrite_existing=overwrite_flag)

src/tools/autofocus_tools.py

@@ -0,0 +1,59 @@
+import cv2
+import numpy as np
+
+def focus_measure_edges(gray: np.ndarray, mask: np.ndarray | None = None) -> float:
+    # mild denoise (optional but usually stabilizes the curve)
+    gray = cv2.GaussianBlur(gray, (3, 3), 0)
+
+    gx = cv2.Scharr(gray, cv2.CV_64F, 1, 0)
+    gy = cv2.Scharr(gray, cv2.CV_64F, 0, 1)
+    g2 = gx * gx + gy * gy
+
+    roi = g2[mask] if mask is not None else g2.reshape(-1)
+    if roi.size == 0:
+        return 0.0
+
+    # threshold relative to median -> knocks out noise floor
+    t = float(np.median(roi) * 3.0)
+    strong = roi[roi > t]
+
+    if strong.size == 0:
+        return 0.0
+
+    print(f"mask pixels: {mask.sum()}, "
+          f"focus={strong.mean():.2f}"
+          f"strong_size={strong.size}")
+
+    return float(strong.mean())  # higher = sharper
+
+def focus_measure_blob_size(gray: np.ndarray, mask: np.ndarray | None = None) -> float:
+    """
+    Measures sharpness for a single bright blob.
+    Higher = sharper (smaller blob).
+    """
+    g = gray.astype(np.float64)
+
+    if mask is not None:
+        g = np.where(mask, g, 0.0)
+
+    # Background subtraction is crucial for blob metrics
+    # Use a large-ish blur as background estimate (tune ksize to your scale)
+    bg = cv2.GaussianBlur(g, (0, 0), sigmaX=10.0, sigmaY=10.0)
+    s = g - bg
+    s[s < 0] = 0.0
+
+    total = float(s.sum())
+    if total <= 0:
+        return 0.0
+
+    h, w = s.shape
+    y, x = np.mgrid[0:h, 0:w]
+
+    cx = float((s * x).sum() / total)
+    cy = float((s * y).sum() / total)
+
+    # intensity-weighted second central moment (variance)
+    var = float((s * ((x - cx) ** 2 + (y - cy) ** 2)).sum() / total)
+
+    # smaller var => sharper, so invert
+    return float(1.0 / (var + 1e-9))

src/tools/file_names_record.py

@@ -0,0 +1,53 @@
+import os
+
+
+def save_file_names_to_record(training_folder, validation_folder, test_folder, output_file):
+    """
+    Reads file names from training, validation, and test folders and saves them to a text file.
+
+    Args:
+        training_folder (str): Path to the training folder.
+        validation_folder (str): Path to the validation folder.
+        test_folder (str): Path to the test folder.
+        output_file (str): Path to the output text file.
+    """
+    try:
+        # Collect file names from the training folder
+        training_files = os.listdir(training_folder)
+        training_files.sort()  # Sort the files alphabetically for consistency
+
+        # Collect file names from the validation folder
+        validation_files = os.listdir(validation_folder)
+        validation_files.sort()
+
+        # Collect file names from the test folder
+        test_files = os.listdir(test_folder)
+        test_files.sort()
+
+        # Save the file names to the output file
+        with open(output_file, 'w') as file:
+            file.write("Training Files:\n")
+            for name in training_files:
+                file.write(name + '\n')
+
+            file.write("\nValidation Files:\n")
+            for name in validation_files:
+                file.write(name + '\n')
+
+            file.write("\nTest Files:\n")
+            for name in test_files:
+                file.write(name + '\n')
+
+        print(f"File names have been successfully saved to {output_file}.")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+# Specify the paths to the folders and output file
+training_folder_path = "/Users/duan_j/Applications/alc/tests/yolo/train/labels"
+validation_folder_path = "/Users/duan_j/Applications/alc/tests/yolo/val/labels"
+test_folder_path = "/Users/duan_j/Applications/alc/tests/yolo/test/labels"
+output_file_path = "/Users/duan_j/Applications/alc/tests/yolo/runs/aare-test/train/img_record.txt"
+
+# Save the file names to the record
+save_file_names_to_record(training_folder_path, validation_folder_path, test_folder_path, output_file_path)

src/tools/generate_splits.py

@@ -0,0 +1,79 @@
+"""
+generate_splits_txt.py
+
+Generates a splits.txt file listing all label filenames across train/test/val
+splits in the format:
+
+    split/labels/<filename>.txt
+
+One entry per line, ordered train → test → val. This file can be consumed by a
+shell script to recreate the directory structure and create symlinks, e.g.:
+
+    while IFS= read -r entry; do
+      split=$(echo "$entry" | cut -d/ -f1)
+      fname=$(basename "$entry" .txt)
+      ln -s "$SOURCE/labels/$fname.txt" "splits/$split/labels/$fname.txt"
+      ln -s "$SOURCE/images/$fname.jpg" "splits/$split/images/$fname.jpg"
+    done < splits.txt
+"""
+
+import argparse
+from pathlib import Path
+
+
+SPLITS = ["train", "test", "val"]
+
+
+def generate_splits_txt(splits_dir: Path, output_file: Path) -> None:
+    entries = []
+    counts = {}
+
+    for split in SPLITS:
+        labels_dir = splits_dir / split / "labels"
+        if not labels_dir.is_dir():
+            print(f"Warning: {labels_dir} not found, skipping.")
+            counts[split] = 0
+            continue
+
+        files = sorted(labels_dir.glob("*.txt"))
+        for f in files:
+            entries.append(f"{split}/labels/{f.name}")
+        counts[split] = len(files)
+
+    output_file.write_text("\n".join(entries) + "\n")
+
+    total = sum(counts.values())
+    print(f"Written {total} entries to {output_file}")
+    for split in SPLITS:
+        print(f"  {split}: {counts.get(split, 0)}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate splits.txt from a splits directory.")
+    parser.add_argument(
+        "--input_dir",
+        type=Path,
+        help="Path to the input folder containing train/val/test splits",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=Path,
+        default=Path(__file__).resolve().parents[2] / "data" / "splits",
+        help="Path to the output folder (default: <project_root>/data/splits)",
+    )
+    args = parser.parse_args()
+
+    splits_dir = args.input_dir
+    output_dir = args.output_dir
+
+    if not splits_dir.is_dir():
+        raise SystemExit(f"Error: input folder not found: {splits_dir}")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_file = output_dir / f"{splits_dir.name}_split.txt"
+
+    generate_splits_txt(splits_dir, output_file)
+
+
+if __name__ == "__main__":
+    main()

src/tools/movetext_images.py

@@ -0,0 +1,29 @@
+import os
+import shutil
+
+
+def get_filename_without_ext(filepath):
+    return os.path.splitext(os.path.basename(filepath))[0]
+
+text_source = "/Users/duan_j/Applications/alc/tests/sample_img/20240109/only_images/C1/labels"
+image_source = "/Users/duan_j/Applications/alc/tests/sample_img/20240109/only_images/C1"
+
+# Get all text files and their base names
+text_files = {get_filename_without_ext(f): f
+              for f in os.listdir(text_source)
+              if f.endswith('.txt')}
+
+# Move matching pairs
+for file in os.listdir(image_source):
+    if file.lower().endswith(('.jpg', '.jpeg')):
+        base_name = get_filename_without_ext(file)
+        if base_name in text_files:
+            # Move image
+            image_source_path = os.path.join(image_source, file)
+            image_dest_path = os.path.join("/Users/duan_j/Applications/alc/tests/yolo/train/images", file)
+            shutil.move(image_source_path, image_dest_path)
+
+            # Move corresponding text file
+            text_source_path = os.path.join(text_source, text_files[base_name])
+            text_dest_path = os.path.join("/Users/duan_j/Applications/alc/tests/yolo/train/labels", text_files[base_name])
+            shutil.move(text_source_path, text_dest_path)

src/tools/pt2engine.py

@@ -0,0 +1,25 @@
+from ultralytics import YOLO
+import torch
+import os
+from ultralytics import settings
+settings.update({"hub": False})
+print("CUDA available:", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("GPU:", torch.cuda.get_device_name(0))
+
+try:
+    import onnx
+    print("ONNX:", onnx.__version__)
+except Exception as e:
+    print("ONNX import failed:", e)
+
+try:
+    import tensorrt as trt
+    print("TensorRT:", trt.__version__)
+except Exception as e:
+    print("TensorRT import failed:", e)
+# Load model
+model = YOLO(f"/opt/ml_dir/AareLC/models/best_yolo26n-seg-overlap-false_2026-04-12.pt")
+
+# Export the model to TensorRT engine format
+model.export(format="engine", imgsz=640, simplify=False, device=0, conf=0.25, iou=0.45)  # Adjust imgsz if needed

src/training/prepare_yolo_seg_dataset_20260306_095355.py

@@ -0,0 +1,344 @@
+#!/usr/bin/env python3
+"""
+Auto-generated by AareLC ML Studio on 2026-03-06T09:54:02.
+Prepare YOLO segmentation train/val folders from source images/labels.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import shutil
+from collections import Counter
+from datetime import datetime, timezone
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+DEFAULT_SOURCE = 'data/img_sets/seg_training'
+DEFAULT_DESTINATION = 'data/img_sets/2026-03-06_seg_training_yolo26n-seg'
+DEFAULT_TRAIN_RATIO = 0.85
+DEFAULT_SEED = 42
+DEFAULT_MODEL_SPEC = 'yolo26n-seg.pt'
+CLASS_NAMES = ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
+
+
+def resolve_user_path(path_value: str) -> Path:
+    p = Path(path_value).expanduser()
+    if not p.is_absolute():
+        p = PROJECT_ROOT / p
+    return p.resolve()
+
+
+def resolve_layout(source_root: Path) -> tuple[Path, Path]:
+    images_dir = source_root / "images"
+    labels_dir = source_root / "labels"
+    if images_dir.is_dir() and labels_dir.is_dir():
+        return images_dir, labels_dir
+    return source_root, source_root
+
+
+def collect_pairs(source_root: Path):
+    images_dir, labels_dir = resolve_layout(source_root)
+    image_files = []
+    for ext in IMAGE_EXTENSIONS:
+        image_files.extend(images_dir.rglob(f"*{ext}"))
+        image_files.extend(images_dir.rglob(f"*{ext.upper()}"))
+    image_files = sorted(set(image_files))
+
+    pairs = []
+    missing_labels = []
+    used_labels = set()
+
+    for img_path in image_files:
+        rel_image = img_path.relative_to(images_dir)
+        label_candidate = (labels_dir / rel_image).with_suffix(".txt")
+        if not label_candidate.exists():
+            flat_fallback = labels_dir / f"{img_path.stem}.txt"
+            if flat_fallback.exists():
+                label_candidate = flat_fallback
+
+        if label_candidate.exists():
+            pairs.append((img_path, label_candidate, rel_image))
+            used_labels.add(label_candidate.resolve())
+        else:
+            missing_labels.append(str(img_path))
+
+    all_labels = {p.resolve() for p in labels_dir.rglob("*.txt")}
+    orphan_labels = sorted(str(p) for p in (all_labels - used_labels))
+    return pairs, missing_labels, orphan_labels, images_dir, labels_dir
+
+
+def split_pairs(pairs, train_ratio: float, seed: int):
+    ordered = list(pairs)
+    rng = random.Random(seed)
+    rng.shuffle(ordered)
+    total = len(ordered)
+    if total == 0:
+        return [], []
+
+    train_count = int(total * train_ratio)
+    if total > 1:
+        train_count = max(1, min(total - 1, train_count))
+    else:
+        train_count = 1
+    return ordered[:train_count], ordered[train_count:]
+
+
+def parse_class_ids(label_path: Path):
+    class_ids = []
+    for line in label_path.read_text(encoding="utf-8").splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        head = line.split()[0]
+        try:
+            class_ids.append(int(float(head)))
+        except ValueError:
+            continue
+    return class_ids
+
+
+def write_dataset_yaml(destination_root: Path):
+    lines = [
+        f"path: {destination_root.resolve()}",
+        "train: train/images",
+        "val: val/images",
+        "",
+        "names:",
+    ]
+    for class_id, class_name in enumerate(CLASS_NAMES):
+        lines.append(f"  {class_id}: {class_name}")
+    dataset_yaml = destination_root / "dataset.yaml"
+    dataset_yaml.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    return dataset_yaml
+
+
+def write_train_script(destination_root: Path, dataset_yaml: Path, model_spec: str):
+    train_script = destination_root / "train_yolo_seg.py"
+    script_text = f"""#!/usr/bin/env python3
+from pathlib import Path
+from ultralytics import YOLO
+import argparse
+
+
+def resolve_model(model_value: str) -> str:
+    p = Path(model_value).expanduser()
+    if p.is_absolute():
+        return str(p)
+    local_candidate = Path(__file__).resolve().parent.parent / p
+    if local_candidate.exists():
+        return str(local_candidate.resolve())
+    return model_value
+
+
+def main():
+    parser = argparse.ArgumentParser(description=\"Train YOLO segmentation model\")
+    parser.add_argument(\"--model\", default='yolo26n-seg.pt', help=\"Local model path or Ultralytics model name\")
+    parser.add_argument(\"--data\", default=None, help=\"Dataset YAML path (default: ./dataset.yaml next to this script)\")
+    parser.add_argument(\"--epochs\", type=int, default=100)
+    parser.add_argument(\"--imgsz\", type=int, default=640)
+    parser.add_argument(\"--batch\", type=int, default=16)
+    parser.add_argument(\"--device\", default=\"0\")
+    parser.add_argument(\"--project\", default=\"runs/segment\")
+    parser.add_argument(\"--name\", default=\"aarelc_seg_train\")
+    args = parser.parse_args()
+
+    model_value = resolve_model(args.model)
+    if args.data:
+        data_yaml = Path(args.data).expanduser().resolve()
+    else:
+        data_yaml = Path(__file__).resolve().parent / "dataset.yaml"
+    if not data_yaml.exists():
+        raise FileNotFoundError(
+            f"Dataset YAML not found: {data_yaml}. "
+            "Run prepare_yolo_seg_dataset first or pass --data /path/to/dataset.yaml."
+        )
+
+    model = YOLO(model_value)
+    model.train(
+        data=str(data_yaml),
+        epochs=args.epochs,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        device=args.device,
+        project=args.project,
+        name=args.name,
+    )
+
+
+if __name__ == \"__main__\":
+    main()
+"""
+    train_script.write_text(script_text, encoding="utf-8")
+    train_script.chmod(0o755)
+    return train_script
+
+
+def copy_pairs(split_name: str, rows, destination_root: Path, dry_run: bool):
+    log_rows = []
+    for image_path, label_path, rel_image in rows:
+        dst_image = destination_root / split_name / "images" / rel_image
+        dst_label = destination_root / split_name / "labels" / rel_image.with_suffix(".txt")
+        log_rows.append(
+            {
+                "source_image": str(image_path),
+                "source_label": str(label_path),
+                "dest_image": str(dst_image),
+                "dest_label": str(dst_label),
+            }
+        )
+        if dry_run:
+            continue
+        dst_image.parent.mkdir(parents=True, exist_ok=True)
+        dst_label.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(image_path, dst_image)
+        shutil.copy2(label_path, dst_label)
+    return log_rows
+
+
+def build_text_log(metadata: dict) -> str:
+    lines = [
+        "YOLO Dataset Preparation Log",
+        "=" * 32,
+        f"Prepared at (UTC): {metadata['prepared_at_utc']}",
+        f"Source root: {metadata['source_root']}",
+        f"Destination root: {metadata['destination_root']}",
+        f"Model spec: {metadata['model_spec']}",
+        f"Training script: {metadata['train_script']}",
+        f"Dry run: {metadata['dry_run']}",
+        f"Train ratio: {metadata['train_ratio']}",
+        f"Seed: {metadata['seed']}",
+        "",
+        f"Total image+label pairs: {metadata['totals']['pairs']}",
+        f"Train pairs: {metadata['totals']['train_pairs']}",
+        f"Validation pairs: {metadata['totals']['val_pairs']}",
+        f"Missing labels: {metadata['totals']['missing_labels']}",
+        f"Orphan labels: {metadata['totals']['orphan_labels']}",
+        "",
+        "Class counts:",
+    ]
+    for class_id, count in sorted(metadata["class_counts"].items(), key=lambda x: int(x[0])):
+        class_name = metadata["class_names"].get(class_id, f"class_{class_id}")
+        lines.append(f"  class {class_id} ({class_name}): {count}")
+
+    lines.append("")
+    lines.append("Train files:")
+    for row in metadata["train_files"]:
+        lines.append(f"  IMG {row['source_image']} -> {row['dest_image']}")
+        lines.append(f"  LBL {row['source_label']} -> {row['dest_label']}")
+
+    lines.append("")
+    lines.append("Validation files:")
+    for row in metadata["val_files"]:
+        lines.append(f"  IMG {row['source_image']} -> {row['dest_image']}")
+        lines.append(f"  LBL {row['source_label']} -> {row['dest_label']}")
+
+    if metadata["missing_labels"]:
+        lines.append("")
+        lines.append("Images missing labels:")
+        for path in metadata["missing_labels"]:
+            lines.append(f"  {path}")
+
+    if metadata["orphan_labels"]:
+        lines.append("")
+        lines.append("Labels without matching images:")
+        for path in metadata["orphan_labels"]:
+            lines.append(f"  {path}")
+
+    return "\n".join(lines) + "\n"
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare YOLO segmentation train/val dataset.")
+    parser.add_argument("--source", default=DEFAULT_SOURCE, help="Source folder with images/labels")
+    parser.add_argument("--destination", default=DEFAULT_DESTINATION, help="Destination dataset root")
+    parser.add_argument("--train-ratio", type=float, default=DEFAULT_TRAIN_RATIO, help="Train split ratio")
+    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="Random seed")
+    parser.add_argument("--dry-run", action="store_true", help="Preview only, do not copy files")
+    args = parser.parse_args()
+
+    if not (0.0 < args.train_ratio < 1.0):
+        raise ValueError("train-ratio must be between 0 and 1")
+
+    source_root = resolve_user_path(args.source)
+    destination_root = resolve_user_path(args.destination)
+    if not source_root.exists():
+        raise FileNotFoundError(f"Source folder not found: {source_root}")
+
+    pairs, missing_labels, orphan_labels, images_dir, labels_dir = collect_pairs(source_root)
+    train_pairs, val_pairs = split_pairs(pairs, args.train_ratio, args.seed)
+
+    class_counts = Counter()
+    for _, label_path, _ in pairs:
+        class_counts.update(parse_class_ids(label_path))
+
+    if not args.dry_run:
+        for split in ("train", "val"):
+            (destination_root / split / "images").mkdir(parents=True, exist_ok=True)
+            (destination_root / split / "labels").mkdir(parents=True, exist_ok=True)
+
+    train_rows = copy_pairs("train", train_pairs, destination_root, args.dry_run)
+    val_rows = copy_pairs("val", val_pairs, destination_root, args.dry_run)
+
+    dataset_yaml = None
+    train_script_path = None
+    if not args.dry_run:
+        dataset_yaml = write_dataset_yaml(destination_root)
+        train_script_path = write_train_script(destination_root, dataset_yaml, DEFAULT_MODEL_SPEC)
+
+    now_utc = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+    class_name_map = {str(i): name for i, name in enumerate(CLASS_NAMES)}
+    for class_id in class_counts:
+        class_name_map.setdefault(str(class_id), f"class_{class_id}")
+
+    metadata = {
+        "prepared_at_utc": now_utc,
+        "source_root": str(source_root),
+        "destination_root": str(destination_root),
+        "images_dir_resolved": str(images_dir),
+        "labels_dir_resolved": str(labels_dir),
+        "dry_run": args.dry_run,
+        "train_ratio": args.train_ratio,
+        "seed": args.seed,
+        "dataset_yaml": str(dataset_yaml) if dataset_yaml else None,
+        "train_script": str(train_script_path) if train_script_path else None,
+        "model_spec": DEFAULT_MODEL_SPEC,
+        "class_names": class_name_map,
+        "class_counts": {str(k): int(v) for k, v in sorted(class_counts.items())},
+        "totals": {
+            "pairs": len(pairs),
+            "train_pairs": len(train_pairs),
+            "val_pairs": len(val_pairs),
+            "missing_labels": len(missing_labels),
+            "orphan_labels": len(orphan_labels),
+        },
+        "train_files": train_rows,
+        "val_files": val_rows,
+        "missing_labels": missing_labels,
+        "orphan_labels": orphan_labels,
+    }
+
+    logs_root = destination_root / "logs"
+    if not args.dry_run:
+        logs_root.mkdir(parents=True, exist_ok=True)
+        stamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
+        json_log = logs_root / f"dataset_preparation_{stamp}.json"
+        txt_log = logs_root / f"dataset_preparation_{stamp}.txt"
+        json_log.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+        txt_log.write_text(build_text_log(metadata), encoding="utf-8")
+        print(f"Wrote logs: {json_log} and {txt_log}")
+        print(f"Wrote training script: {train_script_path}")
+    else:
+        print("Dry run enabled: files were not copied and logs were not written.")
+        print("Summary:")
+        print(json.dumps(metadata["totals"], indent=2))
+
+    print(
+        f"Prepared dataset split from {source_root} -> {destination_root} "
+        f"(train={len(train_pairs)}, val={len(val_pairs)}, total={len(pairs)})."
+    )
+
+
+if __name__ == "__main__":
+    main()

src/training/prepare_yolo_seg_dataset_20260316_120428.py

@@ -0,0 +1,344 @@
+#!/usr/bin/env python3
+"""
+Auto-generated by AareLC ML Studio on 2026-03-16T12:04:37.
+Prepare YOLO segmentation train/val folders from source images/labels.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import shutil
+from collections import Counter
+from datetime import datetime, timezone
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+DEFAULT_SOURCE = ''
+DEFAULT_DESTINATION = 'AareLC/data/img_sets/seg_training_yolo26_20260316'
+DEFAULT_TRAIN_RATIO = 0.85
+DEFAULT_SEED = 42
+DEFAULT_MODEL_SPEC = 'yolo26l-seg.pt'
+CLASS_NAMES = ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
+
+
+def resolve_user_path(path_value: str) -> Path:
+    p = Path(path_value).expanduser()
+    if not p.is_absolute():
+        p = PROJECT_ROOT / p
+    return p.resolve()
+
+
+def resolve_layout(source_root: Path) -> tuple[Path, Path]:
+    images_dir = source_root / "images"
+    labels_dir = source_root / "labels"
+    if images_dir.is_dir() and labels_dir.is_dir():
+        return images_dir, labels_dir
+    return source_root, source_root
+
+
+def collect_pairs(source_root: Path):
+    images_dir, labels_dir = resolve_layout(source_root)
+    image_files = []
+    for ext in IMAGE_EXTENSIONS:
+        image_files.extend(images_dir.rglob(f"*{ext}"))
+        image_files.extend(images_dir.rglob(f"*{ext.upper()}"))
+    image_files = sorted(set(image_files))
+
+    pairs = []
+    missing_labels = []
+    used_labels = set()
+
+    for img_path in image_files:
+        rel_image = img_path.relative_to(images_dir)
+        label_candidate = (labels_dir / rel_image).with_suffix(".txt")
+        if not label_candidate.exists():
+            flat_fallback = labels_dir / f"{img_path.stem}.txt"
+            if flat_fallback.exists():
+                label_candidate = flat_fallback
+
+        if label_candidate.exists():
+            pairs.append((img_path, label_candidate, rel_image))
+            used_labels.add(label_candidate.resolve())
+        else:
+            missing_labels.append(str(img_path))
+
+    all_labels = {p.resolve() for p in labels_dir.rglob("*.txt")}
+    orphan_labels = sorted(str(p) for p in (all_labels - used_labels))
+    return pairs, missing_labels, orphan_labels, images_dir, labels_dir
+
+
+def split_pairs(pairs, train_ratio: float, seed: int):
+    ordered = list(pairs)
+    rng = random.Random(seed)
+    rng.shuffle(ordered)
+    total = len(ordered)
+    if total == 0:
+        return [], []
+
+    train_count = int(total * train_ratio)
+    if total > 1:
+        train_count = max(1, min(total - 1, train_count))
+    else:
+        train_count = 1
+    return ordered[:train_count], ordered[train_count:]
+
+
+def parse_class_ids(label_path: Path):
+    class_ids = []
+    for line in label_path.read_text(encoding="utf-8").splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        head = line.split()[0]
+        try:
+            class_ids.append(int(float(head)))
+        except ValueError:
+            continue
+    return class_ids
+
+
+def write_dataset_yaml(destination_root: Path):
+    lines = [
+        f"path: {destination_root.resolve()}",
+        "train: train/images",
+        "val: val/images",
+        "",
+        "names:",
+    ]
+    for class_id, class_name in enumerate(CLASS_NAMES):
+        lines.append(f"  {class_id}: {class_name}")
+    dataset_yaml = destination_root / "dataset.yaml"
+    dataset_yaml.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    return dataset_yaml
+
+
+def write_train_script(destination_root: Path, dataset_yaml: Path, model_spec: str):
+    train_script = destination_root / "train_yolo_seg.py"
+    script_text = f"""#!/usr/bin/env python3
+from pathlib import Path
+from ultralytics import YOLO
+import argparse
+
+
+def resolve_model(model_value: str) -> str:
+    p = Path(model_value).expanduser()
+    if p.is_absolute():
+        return str(p)
+    local_candidate = Path(__file__).resolve().parent.parent / p
+    if local_candidate.exists():
+        return str(local_candidate.resolve())
+    return model_value
+
+
+def main():
+    parser = argparse.ArgumentParser(description=\"Train YOLO segmentation model\")
+    parser.add_argument(\"--model\", default='yolo26l-seg.pt', help=\"Local model path or Ultralytics model name\")
+    parser.add_argument(\"--data\", default=None, help=\"Dataset YAML path (default: ./dataset.yaml next to this script)\")
+    parser.add_argument(\"--epochs\", type=int, default=100)
+    parser.add_argument(\"--imgsz\", type=int, default=640)
+    parser.add_argument(\"--batch\", type=int, default=16)
+    parser.add_argument(\"--device\", default=\"0\")
+    parser.add_argument(\"--project\", default=\"runs/segment\")
+    parser.add_argument(\"--name\", default=\"aarelc_seg_train\")
+    args = parser.parse_args()
+
+    model_value = resolve_model(args.model)
+    if args.data:
+        dataset_yaml = Path(args.data).expanduser().resolve()
+    else:
+        dataset_yaml = Path(__file__).resolve().parent / "dataset.yaml"
+    if not dataset_yaml.exists():
+        raise FileNotFoundError(
+            f"Dataset YAML not found: {dataset_yaml}. "
+            "Run prepare_yolo_seg_dataset first or pass --data /path/to/dataset.yaml."
+        )
+
+    model = YOLO(model_value)
+    model.train(
+        data=str(dataset_yaml),
+        epochs=args.epochs,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        device=args.device,
+        project=args.project,
+        name=args.name,
+    )
+
+
+if __name__ == \"__main__\":
+    main()
+"""
+    train_script.write_text(script_text, encoding="utf-8")
+    train_script.chmod(0o755)
+    return train_script
+
+
+def copy_pairs(split_name: str, rows, destination_root: Path, dry_run: bool):
+    log_rows = []
+    for image_path, label_path, rel_image in rows:
+        dst_image = destination_root / split_name / "images" / rel_image
+        dst_label = destination_root / split_name / "labels" / rel_image.with_suffix(".txt")
+        log_rows.append(
+            {
+                "source_image": str(image_path),
+                "source_label": str(label_path),
+                "dest_image": str(dst_image),
+                "dest_label": str(dst_label),
+            }
+        )
+        if dry_run:
+            continue
+        dst_image.parent.mkdir(parents=True, exist_ok=True)
+        dst_label.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(image_path, dst_image)
+        shutil.copy2(label_path, dst_label)
+    return log_rows
+
+
+def build_text_log(metadata: dict) -> str:
+    lines = [
+        "YOLO Dataset Preparation Log",
+        "=" * 32,
+        f"Prepared at (UTC): {metadata['prepared_at_utc']}",
+        f"Source root: {metadata['source_root']}",
+        f"Destination root: {metadata['destination_root']}",
+        f"Model spec: {metadata['model_spec']}",
+        f"Training script: {metadata['train_script']}",
+        f"Dry run: {metadata['dry_run']}",
+        f"Train ratio: {metadata['train_ratio']}",
+        f"Seed: {metadata['seed']}",
+        "",
+        f"Total image+label pairs: {metadata['totals']['pairs']}",
+        f"Train pairs: {metadata['totals']['train_pairs']}",
+        f"Validation pairs: {metadata['totals']['val_pairs']}",
+        f"Missing labels: {metadata['totals']['missing_labels']}",
+        f"Orphan labels: {metadata['totals']['orphan_labels']}",
+        "",
+        "Class counts:",
+    ]
+    for class_id, count in sorted(metadata["class_counts"].items(), key=lambda x: int(x[0])):
+        class_name = metadata["class_names"].get(class_id, f"class_{class_id}")
+        lines.append(f"  class {class_id} ({class_name}): {count}")
+
+    lines.append("")
+    lines.append("Train files:")
+    for row in metadata["train_files"]:
+        lines.append(f"  IMG {row['source_image']} -> {row['dest_image']}")
+        lines.append(f"  LBL {row['source_label']} -> {row['dest_label']}")
+
+    lines.append("")
+    lines.append("Validation files:")
+    for row in metadata["val_files"]:
+        lines.append(f"  IMG {row['source_image']} -> {row['dest_image']}")
+        lines.append(f"  LBL {row['source_label']} -> {row['dest_label']}")
+
+    if metadata["missing_labels"]:
+        lines.append("")
+        lines.append("Images missing labels:")
+        for path in metadata["missing_labels"]:
+            lines.append(f"  {path}")
+
+    if metadata["orphan_labels"]:
+        lines.append("")
+        lines.append("Labels without matching images:")
+        for path in metadata["orphan_labels"]:
+            lines.append(f"  {path}")
+
+    return "\n".join(lines) + "\n"
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare YOLO segmentation train/val dataset.")
+    parser.add_argument("--source", default=DEFAULT_SOURCE, help="Source folder with images/labels")
+    parser.add_argument("--destination", default=DEFAULT_DESTINATION, help="Destination dataset root")
+    parser.add_argument("--train-ratio", type=float, default=DEFAULT_TRAIN_RATIO, help="Train split ratio")
+    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="Random seed")
+    parser.add_argument("--dry-run", action="store_true", help="Preview only, do not copy files")
+    args = parser.parse_args()
+
+    if not (0.0 < args.train_ratio < 1.0):
+        raise ValueError("train-ratio must be between 0 and 1")
+
+    source_root = resolve_user_path(args.source)
+    destination_root = resolve_user_path(args.destination)
+    if not source_root.exists():
+        raise FileNotFoundError(f"Source folder not found: {source_root}")
+
+    pairs, missing_labels, orphan_labels, images_dir, labels_dir = collect_pairs(source_root)
+    train_pairs, val_pairs = split_pairs(pairs, args.train_ratio, args.seed)
+
+    class_counts = Counter()
+    for _, label_path, _ in pairs:
+        class_counts.update(parse_class_ids(label_path))
+
+    if not args.dry_run:
+        for split in ("train", "val"):
+            (destination_root / split / "images").mkdir(parents=True, exist_ok=True)
+            (destination_root / split / "labels").mkdir(parents=True, exist_ok=True)
+
+    train_rows = copy_pairs("train", train_pairs, destination_root, args.dry_run)
+    val_rows = copy_pairs("val", val_pairs, destination_root, args.dry_run)
+
+    dataset_yaml = None
+    train_script_path = None
+    if not args.dry_run:
+        dataset_yaml = write_dataset_yaml(destination_root)
+        train_script_path = write_train_script(destination_root, dataset_yaml, DEFAULT_MODEL_SPEC)
+
+    now_utc = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+    class_name_map = {str(i): name for i, name in enumerate(CLASS_NAMES)}
+    for class_id in class_counts:
+        class_name_map.setdefault(str(class_id), f"class_{class_id}")
+
+    metadata = {
+        "prepared_at_utc": now_utc,
+        "source_root": str(source_root),
+        "destination_root": str(destination_root),
+        "images_dir_resolved": str(images_dir),
+        "labels_dir_resolved": str(labels_dir),
+        "dry_run": args.dry_run,
+        "train_ratio": args.train_ratio,
+        "seed": args.seed,
+        "dataset_yaml": str(dataset_yaml) if dataset_yaml else None,
+        "train_script": str(train_script_path) if train_script_path else None,
+        "model_spec": DEFAULT_MODEL_SPEC,
+        "class_names": class_name_map,
+        "class_counts": {str(k): int(v) for k, v in sorted(class_counts.items())},
+        "totals": {
+            "pairs": len(pairs),
+            "train_pairs": len(train_pairs),
+            "val_pairs": len(val_pairs),
+            "missing_labels": len(missing_labels),
+            "orphan_labels": len(orphan_labels),
+        },
+        "train_files": train_rows,
+        "val_files": val_rows,
+        "missing_labels": missing_labels,
+        "orphan_labels": orphan_labels,
+    }
+
+    logs_root = destination_root / "logs"
+    if not args.dry_run:
+        logs_root.mkdir(parents=True, exist_ok=True)
+        stamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
+        json_log = logs_root / f"dataset_preparation_{stamp}.json"
+        txt_log = logs_root / f"dataset_preparation_{stamp}.txt"
+        json_log.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+        txt_log.write_text(build_text_log(metadata), encoding="utf-8")
+        print(f"Wrote logs: {json_log} and {txt_log}")
+        print(f"Wrote training script: {train_script_path}")
+    else:
+        print("Dry run enabled: files were not copied and logs were not written.")
+        print("Summary:")
+        print(json.dumps(metadata["totals"], indent=2))
+
+    print(
+        f"Prepared dataset split from {source_root} -> {destination_root} "
+        f"(train={len(train_pairs)}, val={len(val_pairs)}, total={len(pairs)})."
+    )
+
+
+if __name__ == "__main__":
+    main()

src/training/train_yolo_seg.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+from pathlib import Path
+from ultralytics import YOLO
+import argparse
+
+
+def resolve_model(model_value: str) -> str:
+    p = Path(model_value).expanduser()
+    if p.is_absolute():
+        return str(p)
+    local_candidate = Path(__file__).resolve().parent.parent / p
+    if local_candidate.exists():
+        return str(local_candidate.resolve())
+    return model_value
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train YOLO segmentation model")
+    parser.add_argument("--model", default='/opt/ml_dir/AareLC/models/yolo26n-seg.pt', help="Local model path or Ultralytics model name")
+    parser.add_argument("--data", default=None, help="Dataset YAML path (default: ./dataset.yaml next to this script)")
+    parser.add_argument("--epochs", type=int, default=400)
+    parser.add_argument("--imgsz", type=int, default=640)
+    parser.add_argument("--batch", type=int, default=16)
+    parser.add_argument("--device", default="0")
+    parser.add_argument("--project", default="runs/segment_yolo26n-seg_20260401")
+    parser.add_argument("--name", default="yolo26n-seg_20260401")
+    args = parser.parse_args()
+
+    model_value = resolve_model(args.model)
+    if args.data:
+        dataset_yaml = Path(args.data).expanduser().resolve()
+    else:
+        dataset_yaml = Path(__file__).resolve().parent / "dataset.yaml"
+    if not dataset_yaml.exists():
+        raise FileNotFoundError(
+            f"Dataset YAML not found: /opt/ml_dir/AareLC/data/img_sets/seg_training/AareLC/data/img_sets/seg_training_yolo26_20260316/dataset.yaml. "
+            "Run prepare_yolo_seg_dataset first or pass --data /path/to/dataset.yaml."
+        )
+
+    model = YOLO(model_value)
+    model.train(
+        data=str(dataset_yaml),
+        patience=50,
+        epochs=args.epochs,
+        imgsz=args.imgsz,
+        overlap_mask=False,
+        batch=args.batch,
+        # zoom/size style augmentation
+        scale=0.7,        # stronger zoom in/out effect (try 0.5~0.9)
+        translate=0.15,   # shift object position
+        mosaic=0.5,       # combines scenes at different effective scales
+        close_mosaic=10,  # disable mosaic near end for stability
+        # optional multi-scale training (image size jitter)
+        #multi_scale=True,
+        # keep your flips/geometry
+        fliplr=0.5,
+        flipud=0.3,
+        shear=5.0,
+        perspective=0.0,
+        #workers=0,
+        #device=0,
+        device=(0,1,2,3),
+        project=args.project,
+        name=args.name,
+        amp=False,
+    )
+
+
+if __name__ == "__main__":
+    main()