""" ESP32 camera viewer with drive pad + text terminal (WebSocket commands). Usage: .\\.venv\\Scripts\\python.exe viewer.py --ip 192.168.1.178 Arrow keys / WASD also send directions. Enter sends the text box. Commands go over ws://IP/ws (HTTP kept as fallback). """ from __future__ import annotations import argparse import json import threading import time import tkinter as tk from pathlib import Path from tkinter import colorchooser, filedialog, scrolledtext, ttk import cv2 import numpy as np import requests from PIL import Image, ImageTk from websocket import WebSocketApp from targets import TargetDetector, create_target from targets.orange_ball import METHOD_LABELS, METHODS, OrangeBallDetector from targets.teachable_machine import DEFAULT_MODEL_DIR SAMPLES_DIR = Path(__file__).resolve().parent / "samples" SAMPLE_CLASSES = ("Ball", "Background") class RobotApp: def __init__( self, ip: str, fps: float, timeout: float, detector: TargetDetector | None, ) -> None: self.ip = ip self.base = f"http://{ip}" self.ws_url = f"ws://{ip}/ws" self.period = 1.0 / fps if fps > 0 else 0.0 self.timeout = timeout self.detector = detector self.session = requests.Session() self.session.headers.update( {"Cache-Control": "no-cache", "Pragma": "no-cache"} ) self._stop = threading.Event() self._photo: ImageTk.PhotoImage | None = None self._ws: WebSocketApp | None = None self._ws_ready = threading.Event() self._ws_lock = threading.Lock() self._last_det_log = 0.0 self._raw_frame: np.ndarray | None = None self._frame_lock = threading.Lock() self._preview_scale = 1.0 self._preview_display_size = (640, 480) self._preview_src_size = (640, 480) self.root = tk.Tk() self.root.title(f"linefollower-eye — {ip}") self.root.protocol("WM_DELETE_WINDOW", self.on_close) self.show_mask = tk.BooleanVar(master=self.root, value=False) self.tol_var = tk.IntVar(master=self.root, value=10) self.eyedropper_on = tk.BooleanVar(master=self.root, value=False) self.sample_class = tk.StringVar(master=self.root, value=SAMPLE_CLASSES[0]) self._sample_counts = {name: 0 for name in SAMPLE_CLASSES} for name in SAMPLE_CLASSES: folder = SAMPLES_DIR / name folder.mkdir(parents=True, exist_ok=True) self._sample_counts[name] = sum( 1 for p in folder.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png"} ) main = ttk.Frame(self.root, padding=8) main.grid(row=0, column=0, sticky="nsew") self.root.rowconfigure(0, weight=1) self.root.columnconfigure(0, weight=1) main.rowconfigure(0, weight=1) main.columnconfigure(0, weight=1) self.video = ttk.Label(main, cursor="arrow") self.video.grid(row=0, column=0, sticky="nsew", padx=(0, 8)) self.video.bind("", self._on_video_click) side = ttk.Frame(main) side.grid(row=0, column=1, sticky="ns") ttk.Label(side, text="Drive").grid(row=0, column=0, columnspan=3, pady=(0, 4)) ttk.Button(side, text="↑", width=4, command=lambda: self.send_cmd("up")).grid( row=1, column=1, padx=2, pady=2 ) ttk.Button(side, text="←", width=4, command=lambda: self.send_cmd("left")).grid( row=2, column=0, padx=2, pady=2 ) ttk.Button(side, text="●", width=4, command=lambda: self.send_cmd("stop")).grid( row=2, column=1, padx=2, pady=2 ) ttk.Button(side, text="→", width=4, command=lambda: self.send_cmd("right")).grid( row=2, column=2, padx=2, pady=2 ) ttk.Button(side, text="↓", width=4, command=lambda: self.send_cmd("down")).grid( row=3, column=1, padx=2, pady=2 ) row = 4 if isinstance(self.detector, OrangeBallDetector): row = self._build_tune_panel(side, start_row=row) ttk.Label(side, text="Terminal").grid( row=row, column=0, columnspan=3, sticky="w", pady=(16, 4) ) self.log = scrolledtext.ScrolledText( side, width=36, height=12, state="disabled", wrap="word" ) self.log.grid(row=row + 1, column=0, columnspan=3, sticky="nsew") side.rowconfigure(row + 1, weight=1) entry_row = ttk.Frame(side) entry_row.grid(row=row + 2, column=0, columnspan=3, sticky="ew", pady=(6, 0)) self.entry = ttk.Entry(entry_row) self.entry.pack(side="left", fill="x", expand=True) self.entry.bind("", lambda _e: self.send_text()) ttk.Button(entry_row, text="Send", command=self.send_text).pack( side="left", padx=(6, 0) ) self.root.bind("", lambda _e: self.send_cmd("up")) self.root.bind("", lambda _e: self.send_cmd("down")) self.root.bind("", lambda _e: self.send_cmd("left")) self.root.bind("", lambda _e: self.send_cmd("right")) self.root.bind("", lambda _e: self.send_cmd("stop")) self.root.bind("w", lambda _e: self.send_cmd("up")) self.root.bind("s", lambda _e: self.send_cmd("down")) self.root.bind("a", lambda _e: self.send_cmd("left")) self.root.bind("d", lambda _e: self.send_cmd("right")) self.root.bind("p", self._on_save_sample_key) self.root.bind("P", self._on_save_sample_key) self.append_log(f"HTTP {self.base} WS {self.ws_url}") if self.detector is not None: self.append_log(f"Target: {self.detector.name}") if isinstance(self.detector, OrangeBallDetector): self.append_log("Use Eyedropper, then click the video to sample colour") self.append_log( f"Save training frames: pick class, Save sample (or press P) → {SAMPLES_DIR}" ) else: self.append_log("Target: off") self._ws_thread = threading.Thread(target=self._ws_loop, daemon=True) self._ws_thread.start() self._video_thread = threading.Thread(target=self._video_loop, daemon=True) self._video_thread.start() def _build_tune_panel(self, parent: ttk.Frame, start_row: int) -> int: assert isinstance(self.detector, OrangeBallDetector) h_tol, _, _ = self.detector.get_tolerance() self.tol_var.set(h_tol) ttk.Label(parent, text="Target colour").grid( row=start_row, column=0, columnspan=3, sticky="w", pady=(16, 4) ) start_row += 1 row = ttk.Frame(parent) row.grid(row=start_row, column=0, columnspan=3, sticky="ew", pady=2) self.swatch = tk.Canvas(row, width=36, height=24, highlightthickness=1) self.swatch.pack(side="left") ttk.Button(row, text="Pick colour…", command=self._pick_colour).pack( side="left", padx=(8, 4) ) self.eyedrop_btn = ttk.Checkbutton( row, text="Eyedropper", variable=self.eyedropper_on, command=self._on_eyedropper_toggle, ) self.eyedrop_btn.pack(side="left") start_row += 1 method_row = ttk.Frame(parent) method_row.grid(row=start_row, column=0, columnspan=3, sticky="ew", pady=2) ttk.Label(method_row, text="Method", width=11).pack(side="left") self.method_var = tk.StringVar( master=self.root, value=METHOD_LABELS.get(self.detector.get_method(), "Contour"), ) self.method_combo = ttk.Combobox( method_row, textvariable=self.method_var, values=[METHOD_LABELS[m] for m in METHODS], state="readonly", width=16, ) self.method_combo.pack(side="left", fill="x", expand=True, padx=4) self.method_combo.bind("<>", self._on_method_change) start_row += 1 tm_row = ttk.Frame(parent) tm_row.grid(row=start_row, column=0, columnspan=3, sticky="ew", pady=2) ttk.Button(tm_row, text="Load TM model…", command=self._load_teachable_model).pack( side="left" ) self.tm_status = ttk.Label(tm_row, text="", wraplength=220) self.tm_status.pack(side="left", padx=(8, 0)) self._refresh_tm_status() start_row += 1 ttk.Label(parent, text="Training samples").grid( row=start_row, column=0, columnspan=3, sticky="w", pady=(12, 4) ) start_row += 1 sample_row = ttk.Frame(parent) sample_row.grid(row=start_row, column=0, columnspan=3, sticky="ew", pady=2) ttk.Label(sample_row, text="Class", width=11).pack(side="left") self.sample_combo = ttk.Combobox( sample_row, textvariable=self.sample_class, values=list(SAMPLE_CLASSES), state="readonly", width=14, ) self.sample_combo.pack(side="left", fill="x", expand=True, padx=4) self.sample_combo.bind("<>", lambda _e: self._refresh_sample_count()) start_row += 1 sample_btn_row = ttk.Frame(parent) sample_btn_row.grid(row=start_row, column=0, columnspan=3, sticky="ew", pady=2) ttk.Button(sample_btn_row, text="Save sample", command=self._save_sample).pack( side="left" ) ttk.Button( sample_btn_row, text="Open folder", command=self._open_samples_folder ).pack(side="left", padx=(6, 0)) self.sample_count_label = ttk.Label(sample_btn_row, text="") self.sample_count_label.pack(side="left", padx=(8, 0)) self._refresh_sample_count() start_row += 1 ttk.Checkbutton(parent, text="Show mask", variable=self.show_mask).grid( row=start_row, column=0, columnspan=3, sticky="w" ) start_row += 1 tol_row = ttk.Frame(parent) tol_row.grid(row=start_row, column=0, columnspan=3, sticky="ew", pady=2) ttk.Label(tol_row, text="Tolerance", width=11).pack(side="left") self.tol_label = ttk.Label(tol_row, width=4) ttk.Scale( tol_row, from_=2, to=40, variable=self.tol_var, orient="horizontal", command=self._on_tol_change, ).pack(side="left", fill="x", expand=True, padx=4) self.tol_label.pack(side="left") self._on_tol_change(str(self.tol_var.get())) start_row += 1 self._update_swatch() return start_row def _refresh_sample_count(self) -> None: if not hasattr(self, "sample_count_label"): return name = self.sample_class.get() n = self._sample_counts.get(name, 0) self.sample_count_label.configure(text=f"{name}: {n}") def _on_save_sample_key(self, event: tk.Event) -> str | None: # type: ignore[type-arg] # Don't steal keystrokes from the terminal text box if event.widget is self.entry: return None self._save_sample() return "break" def _save_sample(self) -> None: name = self.sample_class.get().strip() or SAMPLE_CLASSES[0] if name not in SAMPLE_CLASSES: # Allow a custom class typed somehow; still keep it safe as a folder name safe = "".join(ch if ch.isalnum() or ch in "-_ " else "_" for ch in name).strip() name = safe or SAMPLE_CLASSES[0] folder = SAMPLES_DIR / name folder.mkdir(parents=True, exist_ok=True) with self._frame_lock: frame = None if self._raw_frame is None else self._raw_frame.copy() if frame is None: self.append_log("! no frame to save yet") return stamp = time.strftime("%Y%m%d_%H%M%S") millis = int((time.time() % 1) * 1000) path = folder / f"{name}_{stamp}_{millis:03d}.jpg" # cv2.imwrite expects BGR; our raw frame is already BGR ok = cv2.imwrite(str(path), frame, [int(cv2.IMWRITE_JPEG_QUALITY), 95]) if not ok: self.append_log(f"! failed to write {path}") return self._sample_counts[name] = self._sample_counts.get(name, 0) + 1 self._refresh_sample_count() self.append_log(f"saved {path.name} → {folder}") def _open_samples_folder(self) -> None: SAMPLES_DIR.mkdir(parents=True, exist_ok=True) try: import os os.startfile(str(SAMPLES_DIR)) # type: ignore[attr-defined] except Exception as exc: # noqa: BLE001 self.append_log(f"! open folder failed: {exc}") self.append_log(f"samples live in {SAMPLES_DIR}") def _on_method_change(self, _event: object | None = None) -> None: if not isinstance(self.detector, OrangeBallDetector): return label = self.method_var.get() key = next((m for m, text in METHOD_LABELS.items() if text == label), None) if key is None: return self.detector.set_method(key) self.append_log(f"method → {label}") if key == "teachable": self._refresh_tm_status() self.append_log(f"Teachable Machine: {self.detector.teachable_status()}") def _refresh_tm_status(self) -> None: if not isinstance(self.detector, OrangeBallDetector): return if hasattr(self, "tm_status"): self.tm_status.configure(text=self.detector.teachable_status()) def _load_teachable_model(self) -> None: if not isinstance(self.detector, OrangeBallDetector): return initial = self.detector.teachable.model_dir if not initial.is_dir(): initial = DEFAULT_MODEL_DIR path = filedialog.askdirectory( parent=self.root, title="Select Teachable Machine export folder", initialdir=str(initial if initial.is_dir() else Path.cwd()), ) if not path: return self.detector.set_teachable_dir(path) ok = self.detector.teachable.ensure_loaded() self._refresh_tm_status() status = self.detector.teachable_status() self.append_log(f"TM model folder → {path}") self.append_log(f"Teachable Machine: {status}") if ok and self.detector.get_method() != "teachable": self.method_var.set(METHOD_LABELS["teachable"]) self.detector.set_method("teachable") self.append_log("method → Teachable Machine") def _on_eyedropper_toggle(self) -> None: if self.eyedropper_on.get(): self.video.configure(cursor="crosshair") self.append_log("Eyedropper on — click the video") else: self.video.configure(cursor="arrow") def _update_swatch(self) -> None: if not isinstance(self.detector, OrangeBallDetector): return b, g, r = self.detector.get_bgr_center() color = f"#{r:02x}{g:02x}{b:02x}" self.swatch.configure(bg=color) self.swatch.delete("all") self.swatch.create_rectangle(0, 0, 36, 24, fill=color, outline="") def _pick_colour(self) -> None: if not isinstance(self.detector, OrangeBallDetector): return # Don't leave eyedropper armed while dialog is open self.eyedropper_on.set(False) self._on_eyedropper_toggle() b, g, r = self.detector.get_bgr_center() result = colorchooser.askcolor( color=(r, g, b), title="Target ball colour", parent=self.root ) if not result or result[0] is None: return r2, g2, b2 = (int(v) for v in result[0]) self.detector.set_from_bgr(b2, g2, r2) self._update_swatch() self.append_log(f"colour set RGB=({r2},{g2},{b2})") def _on_tol_change(self, _value: str) -> None: if not isinstance(self.detector, OrangeBallDetector): return tol = int(round(float(self.tol_var.get()))) self.tol_label.configure(text=str(tol)) self.detector.set_tolerance(h_tol=tol, s_tol=tol * 6, v_tol=tol * 6) def _video_click_to_frame(self, event: tk.Event) -> tuple[int, int] | None: # type: ignore[type-arg] """Map a click on the video label to pixel coords in the source frame.""" scale = self._preview_scale disp_w, disp_h = self._preview_display_size src_w, src_h = self._preview_src_size if scale <= 0 or disp_w <= 0 or disp_h <= 0: return None # The label widget is often larger than the photo; the image is centred. widget_w = max(1, self.video.winfo_width()) widget_h = max(1, self.video.winfo_height()) offset_x = max(0, (widget_w - disp_w) // 2) offset_y = max(0, (widget_h - disp_h) // 2) ix = event.x - offset_x iy = event.y - offset_y if not (0 <= ix < disp_w and 0 <= iy < disp_h): return None x = int(round(ix / scale)) y = int(round(iy / scale)) x = max(0, min(x, src_w - 1)) y = max(0, min(y, src_h - 1)) return x, y def _on_video_click(self, event: tk.Event) -> None: # type: ignore[type-arg] if not isinstance(self.detector, OrangeBallDetector): return if not self.eyedropper_on.get(): return with self._frame_lock: frame = None if self._raw_frame is None else self._raw_frame.copy() if frame is None: return pt = self._video_click_to_frame(event) if pt is None: self.append_log("! eyedropper: click on the video image") return x, y = pt b, g, r = (int(v) for v in frame[y, x]) self.detector.set_from_bgr(b, g, r) self._update_swatch() self.append_log(f"eyedrop @ ({x},{y}) RGB=({r},{g},{b})") # One-shot: turn off after sampling self.eyedropper_on.set(False) self._on_eyedropper_toggle() def append_log(self, line: str) -> None: def _do() -> None: self.log.configure(state="normal") self.log.insert("end", line + "\n") self.log.see("end") self.log.configure(state="disabled") self.root.after(0, _do) def _ws_loop(self) -> None: while not self._stop.is_set(): self._ws_ready.clear() def on_open(_ws: WebSocketApp) -> None: self._ws_ready.set() self.append_log("WS connected") def on_message(_ws: WebSocketApp, message: str) -> None: self.append_log(f"< {message}") def on_error(_ws: WebSocketApp, error: object) -> None: self.append_log(f"! WS error: {error}") def on_close(_ws: WebSocketApp, *_args: object) -> None: self._ws_ready.clear() self.append_log("WS closed") ws = WebSocketApp( self.ws_url, on_open=on_open, on_message=on_message, on_error=on_error, on_close=on_close, ) with self._ws_lock: self._ws = ws try: ws.run_forever(ping_interval=20, ping_timeout=10) except Exception as exc: # noqa: BLE001 self.append_log(f"! WS run failed: {exc}") with self._ws_lock: self._ws = None if not self._stop.is_set(): time.sleep(1.0) def _ws_send(self, payload: dict) -> bool: data = json.dumps(payload, separators=(",", ":")) with self._ws_lock: ws = self._ws ready = self._ws_ready.is_set() if not ws or not ready: return False try: ws.send(data) return True except Exception as exc: # noqa: BLE001 self.append_log(f"! WS send failed: {exc}") return False def send_cmd(self, direction: str) -> None: def _worker() -> None: if self._ws_send({"type": "cmd", "dir": direction}): return # HTTP fallback try: r = self.session.get( f"{self.base}/cmd", params={"dir": direction}, timeout=self.timeout, ) self.append_log(f"< {r.text.strip()}") except requests.RequestException as exc: self.append_log(f"! cmd {direction} failed: {exc}") threading.Thread(target=_worker, daemon=True).start() def send_text(self) -> None: text = self.entry.get().strip() if not text: return self.entry.delete(0, "end") def _worker() -> None: if self._ws_send({"type": "msg", "text": text}): return try: r = self.session.post( f"{self.base}/msg", data=text.encode("utf-8"), headers={"Content-Type": "text/plain; charset=utf-8"}, timeout=self.timeout, ) self.append_log(f"< {r.text.strip()}") except requests.RequestException as exc: self.append_log(f"! send failed: {exc}") threading.Thread(target=_worker, daemon=True).start() def _show_frame(self, frame_bgr: np.ndarray) -> None: preview_w = 640 h, w = frame_bgr.shape[:2] scale = preview_w / float(w) preview_h = max(1, int(round(h * scale))) self._preview_scale = scale self._preview_display_size = (preview_w, preview_h) self._preview_src_size = (w, h) resized = cv2.resize(frame_bgr, (preview_w, preview_h), interpolation=cv2.INTER_AREA) rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) image = Image.fromarray(rgb) photo = ImageTk.PhotoImage(image=image) self._photo = photo self.video.configure(image=photo) def _video_loop(self) -> None: url = f"{self.base}/capture" while not self._stop.is_set(): t0 = time.perf_counter() try: resp = self.session.get(url, timeout=self.timeout) resp.raise_for_status() buf = np.frombuffer(resp.content, dtype=np.uint8) frame = cv2.imdecode(buf, cv2.IMREAD_COLOR) if frame is not None: with self._frame_lock: self._raw_frame = frame.copy() display = frame if self.detector is not None: if ( isinstance(self.detector, OrangeBallDetector) and self.show_mask.get() ): mask = self.detector.build_mask(frame) display = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) detections = self.detector.detect(frame) display = self.detector.annotate(display, detections) now = time.perf_counter() if detections and now - self._last_det_log > 1.0: top = detections[0] self.append_log( f"{top.label} @ ({top.cx:.0f},{top.cy:.0f}) " f"r={top.radius:.0f} score={top.score:.2f}" ) self._last_det_log = now self.root.after(0, self._show_frame, display) except requests.RequestException as exc: self.append_log(f"! video: {exc}") time.sleep(1.0) if self.period > 0: elapsed = time.perf_counter() - t0 sleep_s = self.period - elapsed if sleep_s > 0: time.sleep(sleep_s) def on_close(self) -> None: self._stop.set() with self._ws_lock: if self._ws is not None: try: self._ws.close() except Exception: # noqa: BLE001 pass self.session.close() self.root.destroy() def run(self) -> None: self.root.mainloop() def main() -> int: parser = argparse.ArgumentParser(description="ESP32 camera + control UI") parser.add_argument("--ip", default="192.168.1.178") parser.add_argument("--fps", type=float, default=10.0) parser.add_argument("--timeout", type=float, default=2.0) parser.add_argument( "--target", default="orange_ball", help="Vision target module: orange_ball | none", ) args = parser.parse_args() detector = create_target(args.target) app = RobotApp( ip=args.ip, fps=args.fps, timeout=args.timeout, detector=detector, ) app.run() return 0 if __name__ == "__main__": raise SystemExit(main())