237 lines
7.9 KiB
Python
237 lines
7.9 KiB
Python
|
"""Teachable Machine (Keras export) helper for ball localisation.
|
||
|
|
|
||
|
|
Export an Image Project from https://teachablemachine.withgoogle.com/
|
||
|
|
as Tensorflow / Keras, then drop into models/teachable/:
|
||
|
|
|
||
|
|
keras_model.h5
|
||
|
|
labels.txt
|
||
|
|
|
||
|
|
TM is a classifier, so we scan a coarse grid of crops and pick the
|
||
|
|
highest-confidence non-background cells as detections.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
import threading
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .base import Detection
|
||
|
|
|
||
|
|
DEFAULT_MODEL_DIR = Path(__file__).resolve().parent.parent / "models" / "teachable"
|
||
|
|
INPUT_SIZE = 224
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_labels(path: Path) -> list[str]:
|
||
|
|
labels: list[str] = []
|
||
|
|
for line in path.read_text(encoding="utf-8").splitlines():
|
||
|
|
line = line.strip()
|
||
|
|
if not line:
|
||
|
|
continue
|
||
|
|
# Teachable Machine lines look like: "0 Ball" or just "Ball"
|
||
|
|
parts = line.split(maxsplit=1)
|
||
|
|
if len(parts) == 2 and parts[0].isdigit():
|
||
|
|
labels.append(parts[1].strip())
|
||
|
|
else:
|
||
|
|
labels.append(line)
|
||
|
|
return labels
|
||
|
|
|
||
|
|
|
||
|
|
def _pick_target_index(labels: list[str]) -> int:
|
||
|
|
"""Prefer a class whose name looks like the object, else class 0."""
|
||
|
|
for i, name in enumerate(labels):
|
||
|
|
low = name.lower()
|
||
|
|
if any(tok in low for tok in ("ball", "target", "object", "positive")):
|
||
|
|
return i
|
||
|
|
for i, name in enumerate(labels):
|
||
|
|
low = name.lower()
|
||
|
|
if "background" in low or "nothing" in low or "empty" in low or "neg" in low:
|
||
|
|
continue
|
||
|
|
return i
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
class TeachableMachineModel:
|
||
|
|
"""Lazy-loaded Teachable Machine Keras classifier + grid search."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model_dir: str | Path | None = None,
|
||
|
|
*,
|
||
|
|
confidence: float = 0.70,
|
||
|
|
grid_cols: int = 5,
|
||
|
|
grid_rows: int = 4,
|
||
|
|
window_scale: float = 0.45,
|
||
|
|
) -> None:
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
self.model_dir = Path(model_dir) if model_dir else DEFAULT_MODEL_DIR
|
||
|
|
self.confidence = float(confidence)
|
||
|
|
self.grid_cols = int(grid_cols)
|
||
|
|
self.grid_rows = int(grid_rows)
|
||
|
|
self.window_scale = float(window_scale)
|
||
|
|
self._model = None
|
||
|
|
self._labels: list[str] = []
|
||
|
|
self._target_index = 0
|
||
|
|
self._load_error: str | None = None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def labels(self) -> list[str]:
|
||
|
|
with self._lock:
|
||
|
|
return list(self._labels)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def load_error(self) -> str | None:
|
||
|
|
with self._lock:
|
||
|
|
return self._load_error
|
||
|
|
|
||
|
|
@property
|
||
|
|
def ready(self) -> bool:
|
||
|
|
with self._lock:
|
||
|
|
return self._model is not None
|
||
|
|
|
||
|
|
def set_model_dir(self, model_dir: str | Path) -> None:
|
||
|
|
with self._lock:
|
||
|
|
self.model_dir = Path(model_dir)
|
||
|
|
self._model = None
|
||
|
|
self._labels = []
|
||
|
|
self._load_error = None
|
||
|
|
|
||
|
|
def ensure_loaded(self) -> bool:
|
||
|
|
with self._lock:
|
||
|
|
if self._model is not None:
|
||
|
|
return True
|
||
|
|
try:
|
||
|
|
self._load_unlocked()
|
||
|
|
return True
|
||
|
|
except Exception as exc: # noqa: BLE001 - surface any load failure in UI
|
||
|
|
self._model = None
|
||
|
|
self._load_error = str(exc)
|
||
|
|
return False
|
||
|
|
|
||
|
|
def _load_unlocked(self) -> None:
|
||
|
|
try:
|
||
|
|
import tensorflow as tf # type: ignore
|
||
|
|
from tensorflow import keras # type: ignore
|
||
|
|
except ImportError as exc:
|
||
|
|
raise RuntimeError(
|
||
|
|
"tensorflow not installed. Run: "
|
||
|
|
"pip install tensorflow (or tensorflow-cpu)"
|
||
|
|
) from exc
|
||
|
|
|
||
|
|
model_path = self.model_dir / "keras_model.h5"
|
||
|
|
labels_path = self.model_dir / "labels.txt"
|
||
|
|
if not model_path.is_file():
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"missing {model_path.name} in {self.model_dir}. "
|
||
|
|
"Export Keras from Teachable Machine and copy it here."
|
||
|
|
)
|
||
|
|
if not labels_path.is_file():
|
||
|
|
raise FileNotFoundError(f"missing {labels_path.name} in {self.model_dir}")
|
||
|
|
|
||
|
|
# Older TM exports pass a legacy `groups=` kwarg into DepthwiseConv2D.
|
||
|
|
class _DepthwiseConv2D(keras.layers.DepthwiseConv2D):
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
kwargs.pop("groups", None)
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
self._model = keras.models.load_model(
|
||
|
|
model_path,
|
||
|
|
compile=False,
|
||
|
|
custom_objects={"DepthwiseConv2D": _DepthwiseConv2D},
|
||
|
|
)
|
||
|
|
self._labels = _parse_labels(labels_path)
|
||
|
|
if not self._labels:
|
||
|
|
raise ValueError(f"no labels found in {labels_path}")
|
||
|
|
self._target_index = _pick_target_index(self._labels)
|
||
|
|
self._load_error = None
|
||
|
|
# Silence TF predict progress spam
|
||
|
|
try:
|
||
|
|
tf.get_logger().setLevel("ERROR")
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def preprocess_bgr_crops(crops_bgr: list[np.ndarray]) -> np.ndarray:
|
||
|
|
"""Teachable Machine normalisation: RGB, 224x224, [-1, 1]."""
|
||
|
|
batch = np.empty((len(crops_bgr), INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32)
|
||
|
|
for i, crop in enumerate(crops_bgr):
|
||
|
|
rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
|
||
|
|
resized = cv2.resize(rgb, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_AREA)
|
||
|
|
batch[i] = (resized.astype(np.float32) / 127.5) - 1.0
|
||
|
|
return batch
|
||
|
|
|
||
|
|
def detect(self, frame_bgr: np.ndarray, max_detections: int = 3) -> list[Detection]:
|
||
|
|
if not self.ensure_loaded():
|
||
|
|
return []
|
||
|
|
|
||
|
|
with self._lock:
|
||
|
|
model = self._model
|
||
|
|
labels = list(self._labels)
|
||
|
|
target_index = self._target_index
|
||
|
|
conf_min = self.confidence
|
||
|
|
cols = self.grid_cols
|
||
|
|
rows = self.grid_rows
|
||
|
|
window_scale = self.window_scale
|
||
|
|
|
||
|
|
assert model is not None
|
||
|
|
h, w = frame_bgr.shape[:2]
|
||
|
|
win = int(round(min(h, w) * window_scale))
|
||
|
|
win = max(48, min(win, h, w))
|
||
|
|
|
||
|
|
crops: list[np.ndarray] = []
|
||
|
|
centres: list[tuple[float, float, float]] = [] # cx, cy, radius
|
||
|
|
for r in range(rows):
|
||
|
|
for c in range(cols):
|
||
|
|
cx = (c + 0.5) * w / cols
|
||
|
|
cy = (r + 0.5) * h / rows
|
||
|
|
x0 = int(round(cx - win / 2))
|
||
|
|
y0 = int(round(cy - win / 2))
|
||
|
|
x0 = max(0, min(x0, w - win))
|
||
|
|
y0 = max(0, min(y0, h - win))
|
||
|
|
crop = frame_bgr[y0 : y0 + win, x0 : x0 + win]
|
||
|
|
if crop.shape[0] < 16 or crop.shape[1] < 16:
|
||
|
|
continue
|
||
|
|
crops.append(crop)
|
||
|
|
centres.append((float(x0 + win / 2), float(y0 + win / 2), float(win * 0.45)))
|
||
|
|
|
||
|
|
if not crops:
|
||
|
|
return []
|
||
|
|
|
||
|
|
batch = self.preprocess_bgr_crops(crops)
|
||
|
|
preds = model.predict(batch, verbose=0)
|
||
|
|
preds = np.asarray(preds)
|
||
|
|
|
||
|
|
found: list[Detection] = []
|
||
|
|
target_name = labels[target_index] if target_index < len(labels) else "target"
|
||
|
|
for i, (cx, cy, radius) in enumerate(centres):
|
||
|
|
scores = preds[i]
|
||
|
|
score = float(scores[target_index])
|
||
|
|
if score < conf_min:
|
||
|
|
continue
|
||
|
|
found.append(
|
||
|
|
Detection(
|
||
|
|
label=f"tm:{target_name}",
|
||
|
|
cx=cx,
|
||
|
|
cy=cy,
|
||
|
|
radius=radius,
|
||
|
|
score=score,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
found.sort(key=lambda d: d.score, reverse=True)
|
||
|
|
# Keep spatially spread detections (simple NMS by distance)
|
||
|
|
kept: list[Detection] = []
|
||
|
|
min_dist = win * 0.55
|
||
|
|
for d in found:
|
||
|
|
if any(
|
||
|
|
(d.cx - k.cx) ** 2 + (d.cy - k.cy) ** 2 < min_dist * min_dist
|
||
|
|
for k in kept
|
||
|
|
):
|
||
|
|
continue
|
||
|
|
kept.append(d)
|
||
|
|
if len(kept) >= max_detections:
|
||
|
|
break
|
||
|
|
return kept
|