vision/pc_vision/viewer.py

646 lines
25 KiB
Python
Raw Normal View History

"""
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("<Button-1>", 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("<Return>", lambda _e: self.send_text())
ttk.Button(entry_row, text="Send", command=self.send_text).pack(
side="left", padx=(6, 0)
)
self.root.bind("<Up>", lambda _e: self.send_cmd("up"))
self.root.bind("<Down>", lambda _e: self.send_cmd("down"))
self.root.bind("<Left>", lambda _e: self.send_cmd("left"))
self.root.bind("<Right>", lambda _e: self.send_cmd("right"))
self.root.bind("<space>", 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("<<ComboboxSelected>>", 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("<<ComboboxSelected>>", 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())