feat: recognition engine and session tracker for multi-client serving
engine.py is the process-wide, thread-safe core the server is built on: one YOLO instance behind a lock, thread-local matcher access, a classification pool so detection is never starved, the stock catalog under its own lock, empty-slot geometry normalized to a reference width (the detector's gap thresholds are absolute pixels), and a small LRU of recent still images for later classify/inspect calls. Detection output is normalized so every client sees the same box shape. tracking.py wraps ultralytics BYTETracker. The important part: ultralytics mints track ids from a process-global counter and resets it when a tracker is constructed, so with two live clients one client's ids would be re-issued while the other still held cached labels. SessionTracker gives each tracker its own counter and never touches the global one. Measured defaults (bench_engine.py, M5): detection p50 26-29 ms on CPU vs 112 ms on MPS with byte-identical boxes, so device defaults to cpu; hybrid classifies at ~51 ms/crop, 92 crops/s across 8 threads. Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
356f7aa502
commit
929d77a802
3 changed files with 1191 additions and 0 deletions
337
recognization/code/bench_engine.py
Normal file
337
recognization/code/bench_engine.py
Normal file
|
|
@ -0,0 +1,337 @@
|
|||
"""
|
||||
Component benchmark for the recognition engine.
|
||||
|
||||
Measures the two things that decide the live-mode defaults: how long detection
|
||||
takes per frame (per device) and how long classification takes per product crop
|
||||
(per method), including parallel throughput.
|
||||
|
||||
Usage (from recognization/):
|
||||
.venv/bin/python code/bench_engine.py
|
||||
.venv/bin/python code/bench_engine.py --devices cpu mps --image image.jpg
|
||||
.venv/bin/python code/bench_engine.py --compare hybrid hybrid_fast
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
project_root = os.path.dirname(current_dir)
|
||||
if os.path.exists(os.path.join(project_root, "stock")):
|
||||
os.chdir(project_root)
|
||||
if current_dir not in sys.path:
|
||||
sys.path.insert(0, current_dir)
|
||||
|
||||
import cv2 # noqa: E402
|
||||
import numpy as np # noqa: E402
|
||||
|
||||
|
||||
def percentile(values: List[float], pct: float) -> float:
|
||||
if not values:
|
||||
return 0.0
|
||||
ordered = sorted(values)
|
||||
index = min(len(ordered) - 1, max(0, int(round((pct / 100.0) * (len(ordered) - 1)))))
|
||||
return ordered[index]
|
||||
|
||||
|
||||
def format_table(rows: List[List[str]], headers: List[str]) -> str:
|
||||
widths = [len(h) for h in headers]
|
||||
for row in rows:
|
||||
for i, cell in enumerate(row):
|
||||
widths[i] = max(widths[i], len(str(cell)))
|
||||
lines = [" | ".join(h.ljust(widths[i]) for i, h in enumerate(headers)),
|
||||
"-+-".join("-" * w for w in widths)]
|
||||
for row in rows:
|
||||
lines.append(" | ".join(str(cell).ljust(widths[i]) for i, cell in enumerate(row)))
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def load_image(path: str, width: Optional[int] = None) -> Any:
|
||||
image = cv2.imread(path)
|
||||
if image is None:
|
||||
raise SystemExit(f"cannot read {path}")
|
||||
if width:
|
||||
scale = width / float(image.shape[1])
|
||||
image = cv2.resize(image, (int(round(image.shape[1] * scale)), int(round(image.shape[0] * scale))))
|
||||
return image
|
||||
|
||||
|
||||
def bench_detection(image_path: str, devices: List[str], runs: int, imgsz: int) -> List[Dict[str, Any]]:
|
||||
from detector import ShelfProductDetector
|
||||
|
||||
weights = "runs/detect/train/weights/best.pt"
|
||||
results: List[Dict[str, Any]] = []
|
||||
native = load_image(image_path)
|
||||
frame = load_image(image_path, width=1280)
|
||||
for device in devices:
|
||||
try:
|
||||
detector = ShelfProductDetector(weights, device=None if device == "default" else device, imgsz=imgsz)
|
||||
except Exception as exc:
|
||||
results.append({"device": device, "error": str(exc)})
|
||||
continue
|
||||
for label, sample in (("native", native), ("1280x720-ish", frame)):
|
||||
latencies: List[float] = []
|
||||
boxes = 0
|
||||
for _ in range(runs):
|
||||
started = time.perf_counter()
|
||||
det = detector.detect(sample)
|
||||
latencies.append((time.perf_counter() - started) * 1000.0)
|
||||
boxes = len(det.get("boxes", []))
|
||||
results.append(
|
||||
{
|
||||
"device": device,
|
||||
"size": label,
|
||||
"shape": f"{sample.shape[1]}x{sample.shape[0]}",
|
||||
"boxes": boxes,
|
||||
"p50_ms": round(statistics.median(latencies), 2),
|
||||
"p95_ms": round(percentile(latencies, 95), 2),
|
||||
"min_ms": round(min(latencies), 2),
|
||||
}
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def box_iou(a: List[int], b: List[int]) -> float:
|
||||
ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
|
||||
ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
|
||||
iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
|
||||
inter = iw * ih
|
||||
if inter <= 0:
|
||||
return 0.0
|
||||
area_a = (a[2] - a[0]) * (a[3] - a[1])
|
||||
area_b = (b[2] - b[0]) * (b[3] - b[1])
|
||||
return inter / max(1, area_a + area_b - inter)
|
||||
|
||||
|
||||
def detection_agreement(a: List[Dict[str, Any]], b: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
if not a or not b:
|
||||
return {"matched": 0, "mean_best_iou": 0.0, "count_ratio": 0.0}
|
||||
ious = []
|
||||
for box_a in a:
|
||||
best = max((box_iou(box_a["box"], box_b["box"]) for box_b in b), default=0.0)
|
||||
ious.append(best)
|
||||
return {
|
||||
"matched": sum(1 for i in ious if i >= 0.5),
|
||||
"mean_best_iou": round(float(np.mean(ious)), 3),
|
||||
"count_ratio": round(len(a) / float(max(1, len(b))), 3),
|
||||
}
|
||||
|
||||
|
||||
def bench_crops(image_path: str, methods: List[str]) -> Dict[str, Any]:
|
||||
from detector import ShelfProductDetector
|
||||
from feature_matcher import ProductFeatureMatcher
|
||||
|
||||
detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
|
||||
matcher = ProductFeatureMatcher(stock_dir="stock")
|
||||
image = load_image(image_path)
|
||||
det = detector.detect(image)
|
||||
crops = []
|
||||
for box in det.get("boxes", []):
|
||||
x1, y1, x2, y2 = box["box"]
|
||||
crops.append(image[y1:y2, x1:x2].copy())
|
||||
report: Dict[str, Any] = {"crop_count": len(crops), "methods": {}}
|
||||
for method in methods:
|
||||
latencies: List[float] = []
|
||||
levels: Dict[str, int] = {}
|
||||
skus: List[str] = []
|
||||
for crop in crops:
|
||||
started = time.perf_counter()
|
||||
result = matcher.classify_crop(crop, method=method)
|
||||
latencies.append((time.perf_counter() - started) * 1000.0)
|
||||
level = result.get("identity_level", "error")
|
||||
levels[level] = levels.get(level, 0) + 1
|
||||
if result.get("best_sku"):
|
||||
skus.append(result["best_sku"])
|
||||
report["methods"][method] = {
|
||||
"p50_ms": round(statistics.median(latencies), 2) if latencies else 0.0,
|
||||
"p95_ms": round(percentile(latencies, 95), 2),
|
||||
"mean_ms": round(float(np.mean(latencies)), 2) if latencies else 0.0,
|
||||
"levels": levels,
|
||||
"accepted_skus": len(set(skus)),
|
||||
}
|
||||
return report
|
||||
|
||||
|
||||
def bench_parallel(crops: List[Any], matcher: Any, method: str, pools: List[int], num_threads: List[int]) -> List[Dict[str, Any]]:
|
||||
rows = []
|
||||
for threads in num_threads:
|
||||
previous = cv2.getNumThreads()
|
||||
cv2.setNumThreads(threads)
|
||||
for pool in pools:
|
||||
started = time.perf_counter()
|
||||
with ThreadPoolExecutor(max_workers=pool) as executor:
|
||||
list(executor.map(lambda c: matcher.classify_crop(c, method=method), crops))
|
||||
elapsed = time.perf_counter() - started
|
||||
rows.append(
|
||||
{
|
||||
"cv2_threads": threads,
|
||||
"pool": pool,
|
||||
"crops": len(crops),
|
||||
"seconds": round(elapsed, 2),
|
||||
"crops_per_s": round(len(crops) / elapsed, 1),
|
||||
}
|
||||
)
|
||||
cv2.setNumThreads(previous)
|
||||
return rows
|
||||
|
||||
|
||||
def compare_methods(image_path: str, methods: List[str], extra_images: Optional[List[str]] = None) -> Dict[str, Any]:
|
||||
"""Agreement between two methods on the same crops (adoption gate)."""
|
||||
from detector import ShelfProductDetector
|
||||
from feature_matcher import ProductFeatureMatcher
|
||||
|
||||
if len(methods) != 2:
|
||||
raise SystemExit("--compare needs exactly two method names")
|
||||
reference, candidate = methods
|
||||
detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
|
||||
matcher = ProductFeatureMatcher(stock_dir="stock")
|
||||
images = [image_path] + (extra_images or [])
|
||||
accepted = 0
|
||||
agree = 0
|
||||
disagree_accepts = 0
|
||||
total = 0
|
||||
for path in images:
|
||||
image = load_image(path)
|
||||
for box in detector.detect(image).get("boxes", []):
|
||||
x1, y1, x2, y2 = box["box"]
|
||||
crop = image[y1:y2, x1:x2].copy()
|
||||
if min(crop.shape[:2]) < 16:
|
||||
continue
|
||||
a = matcher.classify_crop(crop, method=reference)
|
||||
b = matcher.classify_crop(crop, method=candidate)
|
||||
total += 1
|
||||
if a.get("identity_level") == "accepted":
|
||||
accepted += 1
|
||||
if a.get("best_sku") == b.get("best_sku"):
|
||||
agree += 1
|
||||
elif b.get("identity_level") == "accepted":
|
||||
disagree_accepts += 1
|
||||
return {
|
||||
"reference": reference,
|
||||
"candidate": candidate,
|
||||
"crops": total,
|
||||
"reference_accepted": accepted,
|
||||
"agreement_on_accepted": round(agree / accepted, 4) if accepted else None,
|
||||
"conflicting_accepts": disagree_accepts,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Dorayaki engine component benchmark")
|
||||
parser.add_argument("--image", default="image.jpg")
|
||||
parser.add_argument("--devices", nargs="*", default=["default", "cpu", "mps"])
|
||||
parser.add_argument("--runs", type=int, default=15)
|
||||
parser.add_argument("--imgsz", type=int, default=640)
|
||||
parser.add_argument("--methods", nargs="*", default=["surf", "orb", "brisk", "hybrid"])
|
||||
parser.add_argument("--pools", nargs="*", type=int, default=[1, 4, 8])
|
||||
parser.add_argument("--cv2-threads", nargs="*", type=int, default=[0, 1])
|
||||
parser.add_argument("--compare", nargs="*", default=None)
|
||||
parser.add_argument("--compare-images", nargs="*", default=None)
|
||||
parser.add_argument("--json", default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
print("=" * 78)
|
||||
print("DORAYAKI ENGINE BENCHMARK")
|
||||
print("=" * 78)
|
||||
|
||||
from detector import ShelfProductDetector
|
||||
|
||||
started = time.perf_counter()
|
||||
detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
|
||||
load_ms = (time.perf_counter() - started) * 1000.0
|
||||
print(f"model load: {load_ms:.0f} ms")
|
||||
|
||||
warm_image = load_image(args.image, width=1280)
|
||||
started = time.perf_counter()
|
||||
detector.detect(warm_image)
|
||||
print(f"first predict (cold): {(time.perf_counter() - started) * 1000.0:.0f} ms")
|
||||
started = time.perf_counter()
|
||||
detector.detect(warm_image)
|
||||
print(f"second predict (warm): {(time.perf_counter() - started) * 1000.0:.0f} ms")
|
||||
|
||||
print("\n--- detection latency ---")
|
||||
detection = bench_detection(args.image, args.devices, args.runs, args.imgsz)
|
||||
rows = []
|
||||
for row in detection:
|
||||
if "error" in row:
|
||||
rows.append([row["device"], "-", "-", "-", f"ERROR: {row['error'][:40]}"])
|
||||
continue
|
||||
rows.append([row["device"], row["size"], row["shape"], row["boxes"],
|
||||
f"{row['p50_ms']} / {row['p95_ms']}"])
|
||||
print(format_table(rows, ["device", "size", "shape", "boxes", "p50/p95 ms"]))
|
||||
|
||||
# device agreement: do mps boxes match cpu boxes?
|
||||
by_device = {}
|
||||
for device in args.devices:
|
||||
if device in ("default", "cpu"):
|
||||
continue
|
||||
try:
|
||||
detector_d = ShelfProductDetector("runs/detect/train/weights/best.pt",
|
||||
device=None if device == "default" else device)
|
||||
except Exception:
|
||||
continue
|
||||
by_device[device] = detector_d.detect(warm_image).get("boxes", [])
|
||||
cpu_boxes = None
|
||||
try:
|
||||
cpu_detector = ShelfProductDetector("runs/detect/train/weights/best.pt", device="cpu")
|
||||
cpu_boxes = cpu_detector.detect(warm_image).get("boxes", [])
|
||||
except Exception:
|
||||
pass
|
||||
if cpu_boxes is not None:
|
||||
for device, boxes in by_device.items():
|
||||
print(f"agreement cpu vs {device}: {json.dumps(detection_agreement(cpu_boxes, boxes))}")
|
||||
|
||||
print("\n--- classification latency (real crops from", args.image, ") ---")
|
||||
crop_report = bench_crops(args.image, args.methods)
|
||||
print(f"crops: {crop_report['crop_count']}")
|
||||
rows = []
|
||||
for method, stats in crop_report["methods"].items():
|
||||
levels = stats["levels"]
|
||||
rows.append([method, f"{stats['p50_ms']} / {stats['p95_ms']}",
|
||||
f"acc={levels.get('accepted', 0)}",
|
||||
f"amb={levels.get('ambiguous', 0)}",
|
||||
f"unk={levels.get('unknown', 0)}",
|
||||
f"rej={levels.get('rejected', 0)}",
|
||||
f"skus={stats['accepted_skus']}"])
|
||||
print(format_table(rows, ["method", "p50/p95 ms", "accepted", "ambiguous", "unknown", "rejected", "distinct skus"]))
|
||||
|
||||
print("\n--- parallel throughput (hybrid) ---")
|
||||
from feature_matcher import ProductFeatureMatcher
|
||||
|
||||
matcher = ProductFeatureMatcher(stock_dir="stock")
|
||||
image = load_image(args.image)
|
||||
boxes = detector.detect(image).get("boxes", [])
|
||||
crops = [image[b["box"][1]:b["box"][3], b["box"][0]:b["box"][2]].copy() for b in boxes]
|
||||
rows = []
|
||||
for row in bench_parallel(crops, matcher, "hybrid", args.pools, args.cv2_threads):
|
||||
rows.append([row["cv2_threads"], row["pool"], row["crops"], row["seconds"], row["crops_per_s"]])
|
||||
print(format_table(rows, ["cv2_threads", "pool", "crops", "seconds", "crops/s"]))
|
||||
|
||||
if args.compare:
|
||||
print("\n--- method agreement ---")
|
||||
comparison = compare_methods(args.image, args.compare, args.compare_images)
|
||||
print(json.dumps(comparison, indent=2))
|
||||
|
||||
if args.json:
|
||||
with open(args.json, "w") as handle:
|
||||
json.dump(
|
||||
{
|
||||
"model_load_ms": round(load_ms, 1),
|
||||
"detection": detection,
|
||||
"classification": crop_report,
|
||||
},
|
||||
handle,
|
||||
indent=2,
|
||||
)
|
||||
print(f"\nwrote {args.json}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
666
recognization/code/engine.py
Normal file
666
recognization/code/engine.py
Normal file
|
|
@ -0,0 +1,666 @@
|
|||
"""
|
||||
Recognition engine for the Dorayaki shelf-recognition server.
|
||||
|
||||
One process-wide, thread-safe engine. Every call is stateless from the caller's
|
||||
point of view: the active image is never stored here, so any number of clients
|
||||
can call concurrently.
|
||||
|
||||
Owns:
|
||||
- the YOLO product detector (single instance, every predict under a lock),
|
||||
- the feature-matching classifier (thread-local OpenCV extractors, a classify
|
||||
thread pool),
|
||||
- the stock catalog (InventoryManager, mutations under a lock),
|
||||
- empty-slot geometry (resolution-normalized),
|
||||
- a small LRU store of recent still images for later classify/inspect calls.
|
||||
|
||||
Heavy imports (torch, ultralytics, cv2) are pulled in lazily by
|
||||
build_default_engine so tests can inject fakes without loading weights.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import collections
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||
|
||||
# Keep the demo box off the network before ultralytics is imported anywhere.
|
||||
os.environ.setdefault("YOLO_OFFLINE", "1")
|
||||
os.environ.setdefault("YOLO_AUTOINSTALL", "false")
|
||||
|
||||
COORDINATE_CONVENTION = "pixel-xyxy-origin-top-left"
|
||||
REFERENCE_WIDTH = 2048.0
|
||||
MAX_UPSCALE = 2.0
|
||||
DEFAULT_IMAGE_STORE_MAX = 16
|
||||
DEFAULT_IMAGE_STORE_TTL_S = 15 * 60
|
||||
DEFAULT_IMAGE_STORE_MAX_BYTES = 256 * 1024 * 1024
|
||||
|
||||
STATE_DETECTED = "detected"
|
||||
STATE_IDENTIFYING = "identifying"
|
||||
STATE_RECOGNIZED = "recognized"
|
||||
STATE_NEEDS_CONFIRMATION = "needs_confirmation"
|
||||
STATE_UNRECOGNIZED = "unrecognized"
|
||||
|
||||
CLASSIFY_METHODS = ("surf", "orb", "brisk", "hybrid")
|
||||
|
||||
|
||||
class UndecodableImage(ValueError):
|
||||
"""Raised when an uploaded payload is not a readable image."""
|
||||
|
||||
|
||||
class UnknownMethod(ValueError):
|
||||
"""Raised for a classification method outside CLASSIFY_METHODS."""
|
||||
|
||||
|
||||
class ImageExpired(KeyError):
|
||||
"""Raised when an image_id is no longer in the store."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class EngineSettings:
|
||||
weights_path: str = "runs/detect/train/weights/best.pt"
|
||||
stock_dir: str = "stock"
|
||||
stock_csv: str = "stock.csv"
|
||||
# Measured on this project (Apple M5, bench_engine.py): CPU detects in ~27 ms
|
||||
# p50 while MPS takes ~112 ms with identical boxes, so CPU is the default.
|
||||
# Override with --device / DORAYAKI_DEVICE (cpu | mps | cuda | cuda:0).
|
||||
device: Optional[str] = "cpu"
|
||||
imgsz: int = 640
|
||||
classify_workers: int = 0 # 0 -> max(2, cpu_count - 2)
|
||||
live_method: str = "hybrid"
|
||||
live_conf: float = 0.1
|
||||
detect_conf: float = 0.25
|
||||
detect_iou: float = 0.45
|
||||
image_store_max: int = DEFAULT_IMAGE_STORE_MAX
|
||||
image_store_ttl_s: float = DEFAULT_IMAGE_STORE_TTL_S
|
||||
image_store_max_bytes: int = DEFAULT_IMAGE_STORE_MAX_BYTES
|
||||
|
||||
|
||||
@dataclass
|
||||
class StoredImage:
|
||||
image_id: str
|
||||
image: Any # numpy BGR array
|
||||
width: int
|
||||
height: int
|
||||
detections: List[Dict[str, Any]] = field(default_factory=list)
|
||||
created_at: float = 0.0
|
||||
nbytes: int = 0
|
||||
|
||||
|
||||
def identity_state_from_level(level: Optional[str]) -> str:
|
||||
"""Map a decide_identity() level to the wire state vocabulary."""
|
||||
if level == "accepted":
|
||||
return STATE_RECOGNIZED
|
||||
if level in ("ambiguous", "unknown"):
|
||||
return STATE_NEEDS_CONFIRMATION
|
||||
return STATE_UNRECOGNIZED
|
||||
|
||||
|
||||
def identity_payload(decision: Optional[Dict[str, Any]], *, attempts: int = 1) -> Dict[str, Any]:
|
||||
"""One wire-shaped identity object from a decide_identity() result."""
|
||||
if not decision:
|
||||
return {
|
||||
"state": STATE_DETECTED,
|
||||
"sku_id": None,
|
||||
"identity_level": None,
|
||||
"confidence": None,
|
||||
"score": None,
|
||||
"candidates": [],
|
||||
"attempts": 0,
|
||||
"method": None,
|
||||
}
|
||||
level = decision.get("identity_level")
|
||||
state = identity_state_from_level(level)
|
||||
candidates = [
|
||||
{"sku_id": row.get("sku_id"), "score": row.get("score")}
|
||||
for row in (decision.get("ranked_candidates") or [])[:3]
|
||||
]
|
||||
return {
|
||||
"state": state,
|
||||
"sku_id": decision.get("best_sku") if state == STATE_RECOGNIZED else None,
|
||||
"identity_level": level,
|
||||
"confidence": decision.get("confidence"),
|
||||
"score": decision.get("score"),
|
||||
"candidates": candidates,
|
||||
"attempts": attempts,
|
||||
"method": decision.get("method"),
|
||||
}
|
||||
|
||||
|
||||
class ImageStore:
|
||||
"""Bounded LRU of decoded still images keyed by a random image_id."""
|
||||
|
||||
def __init__(self, max_items: int = DEFAULT_IMAGE_STORE_MAX,
|
||||
ttl_s: float = DEFAULT_IMAGE_STORE_TTL_S,
|
||||
max_bytes: int = DEFAULT_IMAGE_STORE_MAX_BYTES):
|
||||
self.max_items = max(1, int(max_items))
|
||||
self.ttl_s = float(ttl_s)
|
||||
self.max_bytes = max(1, int(max_bytes))
|
||||
self._items: "collections.OrderedDict[str, StoredImage]" = collections.OrderedDict()
|
||||
self._bytes = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def put(self, image: Any, detections: Optional[List[Dict[str, Any]]] = None) -> StoredImage:
|
||||
height, width = image.shape[:2]
|
||||
stored = StoredImage(
|
||||
image_id="img_" + os.urandom(8).hex(),
|
||||
image=image,
|
||||
width=int(width),
|
||||
height=int(height),
|
||||
detections=list(detections or []),
|
||||
created_at=time.time(),
|
||||
nbytes=int(image.nbytes),
|
||||
)
|
||||
with self._lock:
|
||||
self._evict_locked()
|
||||
self._items[stored.image_id] = stored
|
||||
self._bytes += stored.nbytes
|
||||
while len(self._items) > self.max_items or self._bytes > self.max_bytes:
|
||||
if len(self._items) <= 1:
|
||||
break
|
||||
_, dropped = self._items.popitem(last=False)
|
||||
self._bytes -= dropped.nbytes
|
||||
return stored
|
||||
|
||||
def get(self, image_id: str) -> StoredImage:
|
||||
with self._lock:
|
||||
stored = self._items.get(image_id)
|
||||
if stored is None or time.time() - stored.created_at > self.ttl_s:
|
||||
if stored is not None:
|
||||
self._items.pop(image_id, None)
|
||||
self._bytes -= stored.nbytes
|
||||
raise ImageExpired(image_id)
|
||||
self._items.move_to_end(image_id)
|
||||
return stored
|
||||
|
||||
def update_detections(self, image_id: str, detections: List[Dict[str, Any]]) -> None:
|
||||
with self._lock:
|
||||
stored = self._items.get(image_id)
|
||||
if stored is not None:
|
||||
stored.detections = list(detections)
|
||||
|
||||
def _evict_locked(self) -> None:
|
||||
now = time.time()
|
||||
for image_id in [k for k, v in self._items.items() if now - v.created_at > self.ttl_s]:
|
||||
dropped = self._items.pop(image_id)
|
||||
self._bytes -= dropped.nbytes
|
||||
|
||||
|
||||
class RecognitionEngine:
|
||||
"""Thread-safe facade over detector + matcher + catalog."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
detector: Any,
|
||||
matcher: Any,
|
||||
inventory: Any,
|
||||
*,
|
||||
settings: Optional[EngineSettings] = None,
|
||||
classify_executor: Optional[ThreadPoolExecutor] = None,
|
||||
owns_executor: bool = True,
|
||||
):
|
||||
self.detector = detector
|
||||
self.matcher = matcher
|
||||
self.inventory = inventory
|
||||
self.settings = settings or EngineSettings()
|
||||
workers = self.settings.classify_workers or max(2, (os.cpu_count() or 4) - 2)
|
||||
self.classify_workers = int(workers)
|
||||
self._executor = classify_executor or ThreadPoolExecutor(
|
||||
max_workers=self.classify_workers, thread_name_prefix="classify"
|
||||
)
|
||||
self._owns_executor = owns_executor
|
||||
self._detect_lock = threading.Lock()
|
||||
self._inventory_lock = threading.Lock()
|
||||
self.image_store = ImageStore(
|
||||
self.settings.image_store_max,
|
||||
self.settings.image_store_ttl_s,
|
||||
self.settings.image_store_max_bytes,
|
||||
)
|
||||
self._catalog_cache: Optional[Tuple[str, List[Dict[str, Any]]]] = None
|
||||
self._started_at = time.time()
|
||||
|
||||
# ------------------------------------------------------------------ lifecycle
|
||||
@property
|
||||
def executor(self) -> ThreadPoolExecutor:
|
||||
return self._executor
|
||||
|
||||
def warm_up(self) -> float:
|
||||
"""One dummy predict so the first real frame is not the slow one."""
|
||||
import numpy as np
|
||||
|
||||
blank = np.zeros((self.settings.imgsz, self.settings.imgsz, 3), dtype="uint8")
|
||||
started = time.perf_counter()
|
||||
self.detect(blank, conf=self.settings.detect_conf)
|
||||
return (time.perf_counter() - started) * 1000.0
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self._owns_executor:
|
||||
self._executor.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
# ------------------------------------------------------------------ images
|
||||
def decode_image(self, payload: bytes) -> Any:
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
if not payload:
|
||||
raise UndecodableImage("empty payload")
|
||||
array = np.frombuffer(payload, dtype=np.uint8)
|
||||
image = cv2.imdecode(array, cv2.IMREAD_COLOR)
|
||||
if image is None or image.size == 0:
|
||||
raise UndecodableImage("payload is not a decodable image")
|
||||
return image
|
||||
|
||||
def encode_jpeg(self, image: Any, quality: int = 85) -> bytes:
|
||||
import cv2
|
||||
|
||||
ok, buffer = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality])
|
||||
if not ok:
|
||||
return b""
|
||||
return bytes(buffer)
|
||||
|
||||
def encode_jpeg_data_url(self, image: Any, quality: int = 85) -> str:
|
||||
payload = self.encode_jpeg(image, quality=quality)
|
||||
if not payload:
|
||||
return ""
|
||||
return "data:image/jpeg;base64," + base64.b64encode(payload).decode("ascii")
|
||||
|
||||
# ------------------------------------------------------------------ detection
|
||||
def detect(self, image: Any, conf: Optional[float] = None) -> List[Dict[str, Any]]:
|
||||
"""Runs the detector under a lock; returns normalized box dicts."""
|
||||
threshold = self.settings.detect_conf if conf is None else float(conf)
|
||||
with self._detect_lock:
|
||||
result = self.detector.detect(
|
||||
image, conf_threshold=threshold, iou_threshold=self.settings.detect_iou
|
||||
)
|
||||
boxes = result.get("boxes") or []
|
||||
normalized: List[Dict[str, Any]] = []
|
||||
for index, box in enumerate(boxes):
|
||||
x1, y1, x2, y2 = [int(v) for v in box["box"]]
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
continue
|
||||
normalized.append(
|
||||
{
|
||||
"id": index,
|
||||
"box": [x1, y1, x2, y2],
|
||||
"confidence": float(box.get("confidence", 0.0)),
|
||||
"class_id": int(box.get("class_id", 0)),
|
||||
"width": x2 - x1,
|
||||
"height": y2 - y1,
|
||||
"center": [(x1 + x2) // 2, (y1 + y2) // 2],
|
||||
}
|
||||
)
|
||||
for index, box in enumerate(normalized):
|
||||
box["id"] = index
|
||||
return normalized
|
||||
|
||||
def crop(self, image: Any, box: Sequence[int]) -> Any:
|
||||
x1, y1, x2, y2 = [int(v) for v in box]
|
||||
height, width = image.shape[:2]
|
||||
x1 = max(0, min(x1, width))
|
||||
x2 = max(0, min(x2, width))
|
||||
y1 = max(0, min(y1, height))
|
||||
y2 = max(0, min(y2, height))
|
||||
return image[y1:y2, x1:x2].copy()
|
||||
|
||||
# ------------------------------------------------------------------ empty slots
|
||||
def empty_slots(self, image_shape: Sequence[int], boxes: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Empty-shelf gaps, computed in a 2048 px-wide reference frame so the
|
||||
detector's absolute pixel thresholds do not change with capture size."""
|
||||
height, width = int(image_shape[0]), int(image_shape[1])
|
||||
if not boxes or width <= 0 or height <= 0:
|
||||
return []
|
||||
scale = min(MAX_UPSCALE, REFERENCE_WIDTH / float(width))
|
||||
ref_boxes = []
|
||||
for box in boxes:
|
||||
x1, y1, x2, y2 = [int(v) for v in box["box"]]
|
||||
ref_boxes.append(
|
||||
{
|
||||
"box": [
|
||||
int(round(x1 * scale)),
|
||||
int(round(y1 * scale)),
|
||||
int(round(x2 * scale)),
|
||||
int(round(y2 * scale)),
|
||||
]
|
||||
}
|
||||
)
|
||||
ref_shape = (int(round(height * scale)), int(round(width * scale)))
|
||||
slots = self.detector.detect_empty_slots(ref_shape, ref_boxes)
|
||||
if scale == 1.0:
|
||||
return slots
|
||||
inv = 1.0 / scale
|
||||
for slot in slots:
|
||||
x1, y1, x2, y2 = slot["box"]
|
||||
slot["box"] = [int(round(x1 * inv)), int(round(y1 * inv)),
|
||||
int(round(x2 * inv)), int(round(y2 * inv))]
|
||||
slot["width"] = int(round(slot["width"] * inv))
|
||||
slot["height"] = int(round(slot["height"] * inv))
|
||||
slot["center"] = [int(round(slot["center"][0] * inv)), int(round(slot["center"][1] * inv))]
|
||||
return slots
|
||||
|
||||
# ------------------------------------------------------------------ classification
|
||||
def score_crop(self, crop: Any, method: Optional[str] = None) -> Dict[str, float]:
|
||||
method = (method or self.settings.live_method).lower()
|
||||
if method not in CLASSIFY_METHODS:
|
||||
raise UnknownMethod(method)
|
||||
scores, _ = self.matcher.score_crop(crop, method=method)
|
||||
return {str(sku): float(score) for sku, score in scores.items()}
|
||||
|
||||
def classify_boxes(
|
||||
self,
|
||||
image: Any,
|
||||
boxes: Sequence[Dict[str, Any]],
|
||||
method: Optional[str] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Classifies every box in parallel; returns one entry per box, ordered by area."""
|
||||
from identity import decide_identity
|
||||
|
||||
method = (method or self.settings.live_method).lower()
|
||||
if method not in CLASSIFY_METHODS:
|
||||
raise UnknownMethod(method)
|
||||
if not boxes:
|
||||
return []
|
||||
ordered = sorted(boxes, key=lambda b: b["width"] * b["height"], reverse=True)
|
||||
crops = [self.crop(image, box["box"]) for box in ordered]
|
||||
futures = [self._executor.submit(self.matcher.score_crop, crop, method) for crop in crops]
|
||||
results: List[Dict[str, Any]] = []
|
||||
for box, future in zip(ordered, futures):
|
||||
try:
|
||||
scores, details = future.result()
|
||||
except Exception as exc: # a single bad crop must not fail the batch
|
||||
scores, details = {}, {"error": str(exc)}
|
||||
decision = decide_identity(scores, details=details, method=method)
|
||||
results.append(
|
||||
{
|
||||
"detection_id": int(box["id"]),
|
||||
"box": list(box["box"]),
|
||||
"identity": identity_payload(decision),
|
||||
}
|
||||
)
|
||||
return results
|
||||
|
||||
def inspect(self, image: Any, box: Sequence[int], method: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Full single-crop result: decision, per-method details, crop, visualization."""
|
||||
from identity import decide_identity
|
||||
|
||||
method = (method or self.settings.live_method).lower()
|
||||
if method not in CLASSIFY_METHODS:
|
||||
raise UnknownMethod(method)
|
||||
crop = self.crop(image, box)
|
||||
scores, details = self.matcher.score_crop(crop, method=method)
|
||||
decision = decide_identity(scores, details=details, method=method)
|
||||
identity = identity_payload(decision)
|
||||
|
||||
top_details = decision.get("details") or {}
|
||||
sub_details = None
|
||||
if isinstance(top_details, dict):
|
||||
raw = top_details.get("sub_details")
|
||||
if raw:
|
||||
sub_details = {k: dict(v) for k, v in raw.items()}
|
||||
|
||||
product = None
|
||||
visualization = None
|
||||
visualized_sku = None
|
||||
is_candidate = False
|
||||
if identity["state"] == STATE_RECOGNIZED:
|
||||
product = self.product(identity["sku_id"])
|
||||
visualized_sku = identity["sku_id"]
|
||||
elif identity["candidates"]:
|
||||
visualized_sku = identity["candidates"][0]["sku_id"]
|
||||
is_candidate = True
|
||||
if visualized_sku:
|
||||
try:
|
||||
visualization = self.matcher.generate_match_visualization(
|
||||
crop, visualized_sku, method=method
|
||||
)
|
||||
except Exception:
|
||||
visualization = None
|
||||
|
||||
return {
|
||||
"box": [int(v) for v in box],
|
||||
"method": method,
|
||||
"identity": identity,
|
||||
"sub_details": sub_details,
|
||||
"product": product,
|
||||
"visualized_sku_id": visualized_sku,
|
||||
"is_candidate": is_candidate,
|
||||
"crop_jpeg_b64": self.encode_jpeg_data_url(crop),
|
||||
"match_visualization_jpeg_b64": (
|
||||
self.encode_jpeg_data_url(visualization) if visualization is not None else None
|
||||
),
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------ catalog
|
||||
def catalog_version(self) -> str:
|
||||
digest = hashlib.sha256()
|
||||
csv_path = os.path.join(self.settings.stock_dir, self.settings.stock_csv)
|
||||
try:
|
||||
with open(csv_path, "rb") as handle:
|
||||
digest.update(handle.read())
|
||||
except OSError:
|
||||
digest.update(b"missing-csv")
|
||||
try:
|
||||
entries = sorted(
|
||||
(name, os.path.getsize(os.path.join(self.settings.stock_dir, name)))
|
||||
for name in os.listdir(self.settings.stock_dir)
|
||||
if name.lower().endswith((".jpeg", ".jpg", ".png"))
|
||||
)
|
||||
except OSError:
|
||||
entries = []
|
||||
for name, size in entries:
|
||||
digest.update(f"{name}:{size}".encode("utf-8"))
|
||||
return digest.hexdigest()[:12]
|
||||
|
||||
def model_version(self) -> str:
|
||||
try:
|
||||
import ultralytics
|
||||
|
||||
library = f"ultralytics-{ultralytics.__version__}"
|
||||
except Exception:
|
||||
library = "ultralytics-unknown"
|
||||
return f"{os.path.basename(self.settings.weights_path)}@{library}"
|
||||
|
||||
def catalog(self) -> List[Dict[str, Any]]:
|
||||
version = self.catalog_version()
|
||||
cached = self._catalog_cache
|
||||
if cached and cached[0] == version:
|
||||
return cached[1]
|
||||
products: List[Dict[str, Any]] = []
|
||||
for sku_id in sorted(self.inventory.list_all_skus()):
|
||||
product = self.inventory.get_product(sku_id) or {}
|
||||
products.append(self._product_payload(sku_id, product))
|
||||
self._catalog_cache = (version, products)
|
||||
return products
|
||||
|
||||
def product(self, sku_id: Optional[str]) -> Optional[Dict[str, Any]]:
|
||||
if not sku_id:
|
||||
return None
|
||||
with self._inventory_lock:
|
||||
product = self.inventory.get_product(sku_id)
|
||||
if product is None:
|
||||
return None
|
||||
payload = self._product_payload(sku_id, product)
|
||||
self._catalog_cache = None
|
||||
return payload
|
||||
|
||||
def _product_payload(self, sku_id: str, product: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {
|
||||
"sku_id": sku_id,
|
||||
"product_name": product.get("product_name", ""),
|
||||
"price": product.get("price", 0.0),
|
||||
"number_in": int(product.get("number_in", 0)),
|
||||
"number_out": int(product.get("number_out", 0)),
|
||||
"number_on": int(product.get("number_on", 0)),
|
||||
"receipt_date": product.get("receipt_date", ""),
|
||||
"expiration_date": product.get("expiration_date", ""),
|
||||
"image_url": f"/v1/catalog/{sku_id}/image",
|
||||
}
|
||||
|
||||
def catalog_image_path(self, sku_id: str) -> Optional[str]:
|
||||
"""Resolve a template image path through the inventory map only."""
|
||||
with self._inventory_lock:
|
||||
product = self.inventory.get_product(sku_id)
|
||||
if product is None:
|
||||
return None
|
||||
path = product.get("local_image_path")
|
||||
if not path or not os.path.exists(path):
|
||||
return None
|
||||
return path
|
||||
|
||||
def update_stock(self, sku_id: str, delta: int) -> Optional[Dict[str, Any]]:
|
||||
with self._inventory_lock:
|
||||
ok = self.inventory.update_stock(sku_id, int(delta))
|
||||
if not ok:
|
||||
return None
|
||||
product = self.inventory.get_product(sku_id) or {}
|
||||
payload = self._product_payload(sku_id, product)
|
||||
self._catalog_cache = None
|
||||
return payload
|
||||
|
||||
def sku_count(self) -> int:
|
||||
return len(self.inventory.list_all_skus())
|
||||
|
||||
def uptime_s(self) -> float:
|
||||
return time.time() - self._started_at
|
||||
|
||||
# ------------------------------------------------------------------ recognize
|
||||
def recognize(
|
||||
self,
|
||||
image: Any,
|
||||
*,
|
||||
classify: bool = True,
|
||||
method: Optional[str] = None,
|
||||
conf: Optional[float] = None,
|
||||
with_empty_slots: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""Stateless still-image recognition (photo mode and one-shot scans)."""
|
||||
started = time.perf_counter()
|
||||
boxes = self.detect(image, conf=conf)
|
||||
detected_at = time.perf_counter()
|
||||
detections: List[Dict[str, Any]] = []
|
||||
identities: List[Dict[str, Any]] = []
|
||||
if classify and boxes:
|
||||
identities = self.classify_boxes(image, boxes, method=method)
|
||||
identities_by_id = {row["detection_id"]: row["identity"] for row in identities}
|
||||
for box in boxes:
|
||||
detections.append(
|
||||
{
|
||||
"id": box["id"],
|
||||
"box": box["box"],
|
||||
"confidence": box["confidence"],
|
||||
"identity": identities_by_id.get(box["id"], identity_payload(None)),
|
||||
}
|
||||
)
|
||||
slots = self.empty_slots(image.shape, boxes) if with_empty_slots else []
|
||||
classified_at = time.perf_counter()
|
||||
summary = self.summarize(detections, slots)
|
||||
stored = self.image_store.put(image, detections)
|
||||
return {
|
||||
"image_id": stored.image_id,
|
||||
"width": stored.width,
|
||||
"height": stored.height,
|
||||
"catalog_version": self.catalog_version(),
|
||||
"model_version": self.model_version(),
|
||||
"coordinate_convention": COORDINATE_CONVENTION,
|
||||
"detections": detections,
|
||||
"empty_slots": slots,
|
||||
"summary": summary,
|
||||
"timing_ms": {
|
||||
"detect": round((detected_at - started) * 1000.0, 2),
|
||||
"classify": round((classified_at - detected_at) * 1000.0, 2),
|
||||
"total": round((classified_at - started) * 1000.0, 2),
|
||||
},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def summarize(
|
||||
detections: Sequence[Dict[str, Any]],
|
||||
slots: Sequence[Dict[str, Any]],
|
||||
*,
|
||||
include_identifying: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
counts = {
|
||||
STATE_RECOGNIZED: 0,
|
||||
STATE_NEEDS_CONFIRMATION: 0,
|
||||
STATE_UNRECOGNIZED: 0,
|
||||
STATE_DETECTED: 0,
|
||||
}
|
||||
if include_identifying:
|
||||
counts[STATE_IDENTIFYING] = 0
|
||||
facings: Dict[str, int] = {}
|
||||
for row in detections:
|
||||
state = (row.get("identity") or {}).get("state") or STATE_DETECTED
|
||||
if state not in counts:
|
||||
state = STATE_DETECTED
|
||||
counts[state] += 1
|
||||
sku_id = (row.get("identity") or {}).get("sku_id")
|
||||
if state == STATE_RECOGNIZED and sku_id:
|
||||
facings[sku_id] = facings.get(sku_id, 0) + 1
|
||||
summary: Dict[str, Any] = {
|
||||
"in_view": len(detections),
|
||||
"recognized": counts[STATE_RECOGNIZED],
|
||||
"needs_confirmation": counts[STATE_NEEDS_CONFIRMATION],
|
||||
"unrecognized": counts[STATE_UNRECOGNIZED],
|
||||
"detected": counts[STATE_DETECTED],
|
||||
"empty_slots": len(slots),
|
||||
"by_sku": [
|
||||
{"sku_id": sku, "facings": count}
|
||||
for sku, count in sorted(facings.items(), key=lambda kv: (-kv[1], kv[0]))
|
||||
],
|
||||
}
|
||||
if include_identifying:
|
||||
summary["identifying"] = counts[STATE_IDENTIFYING]
|
||||
return summary
|
||||
|
||||
|
||||
def build_default_engine(settings: Optional[EngineSettings] = None) -> RecognitionEngine:
|
||||
"""Construct the real detector/matcher/inventory. Heavy imports happen here."""
|
||||
import sys
|
||||
|
||||
settings = settings or EngineSettings()
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if current_dir not in sys.path:
|
||||
sys.path.insert(0, current_dir)
|
||||
|
||||
from detector import ShelfProductDetector
|
||||
from feature_matcher import ProductFeatureMatcher
|
||||
from inventory import InventoryManager
|
||||
|
||||
detector = ShelfProductDetector(
|
||||
weights_path=settings.weights_path, device=settings.device, imgsz=settings.imgsz
|
||||
)
|
||||
matcher = ProductFeatureMatcher(stock_dir=settings.stock_dir)
|
||||
inventory = InventoryManager(stock_dir=settings.stock_dir, csv_filename=settings.stock_csv)
|
||||
return RecognitionEngine(detector, matcher, inventory, settings=settings)
|
||||
|
||||
|
||||
def settings_from_env(env: Optional[Dict[str, str]] = None) -> EngineSettings:
|
||||
"""Build EngineSettings from DORAYAKI_* environment variables."""
|
||||
env = dict(os.environ if env is None else env)
|
||||
base = EngineSettings()
|
||||
return replace(
|
||||
base,
|
||||
weights_path=env.get("DORAYAKI_WEIGHTS", base.weights_path),
|
||||
stock_dir=env.get("DORAYAKI_STOCK_DIR", base.stock_dir),
|
||||
device=env.get("DORAYAKI_DEVICE") or None,
|
||||
imgsz=int(env.get("DORAYAKI_IMGSZ", base.imgsz)),
|
||||
classify_workers=int(env.get("DORAYAKI_CLASSIFY_WORKERS", base.classify_workers)),
|
||||
live_method=env.get("DORAYAKI_LIVE_METHOD", base.live_method),
|
||||
live_conf=float(env.get("DORAYAKI_LIVE_CONF", base.live_conf)),
|
||||
detect_conf=float(env.get("DORAYAKI_DETECT_CONF", base.detect_conf)),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__": # tiny smoke check: python code/engine.py
|
||||
engine = build_default_engine(settings_from_env())
|
||||
print("[Engine] warmup ms:", round(engine.warm_up(), 1))
|
||||
print("[Engine] catalog:", engine.sku_count(), "skus, version", engine.catalog_version())
|
||||
result = engine.recognize(engine.decode_image(open("image.jpg", "rb").read()))
|
||||
print("[Engine] detected:", len(result["detections"]), "recognized:", result["summary"]["recognized"])
|
||||
print("[Engine] timing:", json.dumps(result["timing_ms"]))
|
||||
engine.shutdown()
|
||||
188
recognization/code/tracking.py
Normal file
188
recognization/code/tracking.py
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
"""
|
||||
Session-scoped multi-object tracking for live recognition.
|
||||
|
||||
Wraps ultralytics' BYTETracker (Kalman prediction + two-stage association) with
|
||||
one critical fix: ultralytics mints track IDs from a PROCESS-GLOBAL counter
|
||||
(BaseTrack._count), and BYTETracker.__init__ resets that counter. With several
|
||||
live sessions that means one client's tracker creation re-issues IDs already in
|
||||
use by another client, and a cached label can land on the wrong product.
|
||||
|
||||
SessionTracker gives every tracker instance its own counter and never touches
|
||||
the global one. The rest of the algorithm is used unmodified.
|
||||
|
||||
Pinned to ultralytics 8.4.153 (see requirements.txt): this module touches
|
||||
tracker internals.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from typing import Any, List, Optional, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if current_dir not in sys.path:
|
||||
sys.path.insert(0, current_dir)
|
||||
|
||||
from ultralytics.trackers.basetrack import BaseTrack, STrack
|
||||
from ultralytics.trackers.byte_tracker import BYTETracker
|
||||
from ultralytics.utils import IterableSimpleNamespace, yaml_load
|
||||
|
||||
_BYTETRACK_YAML = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__import__("ultralytics").__file__)),
|
||||
"cfg", "trackers", "bytetrack.yaml",
|
||||
)
|
||||
|
||||
DEFAULT_TRACKER_CONFIG = {
|
||||
"tracker_type": "bytetrack",
|
||||
"track_high_thresh": 0.25,
|
||||
"track_low_thresh": 0.1,
|
||||
"new_track_thresh": 0.25,
|
||||
"track_buffer": 30,
|
||||
"match_thresh": 0.8,
|
||||
"fuse_score": True,
|
||||
}
|
||||
|
||||
|
||||
def _default_args() -> IterableSimpleNamespace:
|
||||
try:
|
||||
return IterableSimpleNamespace(**yaml_load(_BYTETRACK_YAML))
|
||||
except Exception:
|
||||
return IterableSimpleNamespace(**DEFAULT_TRACKER_CONFIG)
|
||||
|
||||
|
||||
class _Detections:
|
||||
"""Duck-typed stand-in for ultralytics Results: only what BYTETracker reads.
|
||||
|
||||
Must NOT expose `xywhr`; the tracker uses its presence to pick rotated-box
|
||||
handling.
|
||||
"""
|
||||
|
||||
__slots__ = ("xywh", "conf", "cls")
|
||||
|
||||
def __init__(self, boxes_xyxy: np.ndarray, scores: np.ndarray):
|
||||
boxes = np.asarray(boxes_xyxy, dtype=np.float32).reshape(-1, 4)
|
||||
scores = np.asarray(scores, dtype=np.float32).reshape(-1)
|
||||
widths = boxes[:, 2] - boxes[:, 0]
|
||||
heights = boxes[:, 3] - boxes[:, 1]
|
||||
self.xywh = np.stack(
|
||||
[boxes[:, 0] + widths / 2.0, boxes[:, 1] + heights / 2.0, widths, heights], axis=1
|
||||
).astype(np.float32)
|
||||
self.conf = scores
|
||||
self.cls = np.zeros(len(scores), dtype=np.float32)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.conf)
|
||||
|
||||
def __getitem__(self, mask: Any) -> "_Detections":
|
||||
mask = np.asarray(mask)
|
||||
subset = _Detections.__new__(_Detections)
|
||||
subset.xywh = self.xywh[mask]
|
||||
subset.conf = self.conf[mask]
|
||||
subset.cls = self.cls[mask]
|
||||
return subset
|
||||
|
||||
|
||||
class _SessionByteTrack(BYTETracker):
|
||||
"""BYTETracker with a per-instance ID counter."""
|
||||
|
||||
def __init__(self, args: IterableSimpleNamespace):
|
||||
self._id_counter = itertools.count(1)
|
||||
self._id_lock = threading.Lock()
|
||||
# super().__init__ calls self.reset_id(); the override below only touches
|
||||
# this instance's counter, so the global one is never disturbed.
|
||||
super().__init__(args)
|
||||
self.track_class = self._make_track_class()
|
||||
|
||||
def _next_id(self) -> int:
|
||||
with self._id_lock:
|
||||
return next(self._id_counter)
|
||||
|
||||
def _make_track_class(self):
|
||||
tracker = self
|
||||
|
||||
class SessionSTrack(STrack):
|
||||
@staticmethod
|
||||
def next_id() -> int:
|
||||
return tracker._next_id()
|
||||
|
||||
SessionSTrack.__name__ = "SessionSTrack"
|
||||
return SessionSTrack
|
||||
|
||||
def reset_id(self) -> None: # type: ignore[override]
|
||||
"""Reset only this tracker's counter; never BaseTrack._count."""
|
||||
with self._id_lock:
|
||||
self._id_counter = itertools.count(1)
|
||||
self.track_class = self._make_track_class()
|
||||
|
||||
|
||||
class TrackOut:
|
||||
"""One tracked object for a frame."""
|
||||
|
||||
__slots__ = ("track_id", "box", "score", "det_index")
|
||||
|
||||
def __init__(self, track_id: int, box: List[int], score: float, det_index: Optional[int]):
|
||||
self.track_id = int(track_id)
|
||||
self.box = [int(v) for v in box]
|
||||
self.score = float(score)
|
||||
self.det_index = None if det_index is None or det_index < 0 else int(det_index)
|
||||
|
||||
def __repr__(self) -> str: # pragma: no cover - debug helper
|
||||
return f"TrackOut(id={self.track_id}, box={self.box}, score={self.score:.3f}, det={self.det_index})"
|
||||
|
||||
|
||||
class SessionTracker:
|
||||
"""Per-session tracker. Used by one session at a time, under its lock."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
high: float = 0.25,
|
||||
low: float = 0.1,
|
||||
new: float = 0.25,
|
||||
buffer: int = 30,
|
||||
match: float = 0.8,
|
||||
fuse: bool = True,
|
||||
):
|
||||
args = _default_args()
|
||||
args.track_high_thresh = float(high)
|
||||
args.track_low_thresh = float(low)
|
||||
args.new_track_thresh = float(new)
|
||||
args.track_buffer = int(buffer)
|
||||
args.match_thresh = float(match)
|
||||
args.fuse_score = bool(fuse)
|
||||
self.args = args
|
||||
self._tracker = _SessionByteTrack(args)
|
||||
|
||||
def update(
|
||||
self,
|
||||
boxes_xyxy: Sequence[Sequence[int]],
|
||||
scores: Sequence[float],
|
||||
) -> List[TrackOut]:
|
||||
boxes = np.asarray(boxes_xyxy, dtype=np.float32).reshape(-1, 4)
|
||||
scores_array = np.asarray(scores, dtype=np.float32).reshape(-1)
|
||||
if len(boxes) != len(scores_array):
|
||||
raise ValueError("boxes and scores must have the same length")
|
||||
detections = _Detections(boxes, scores_array)
|
||||
rows = self._tracker.update(detections)
|
||||
outputs: List[TrackOut] = []
|
||||
if rows is None or len(rows) == 0:
|
||||
return outputs
|
||||
for row in np.asarray(rows):
|
||||
# [x1, y1, x2, y2, track_id, score, cls, idx]
|
||||
outputs.append(
|
||||
TrackOut(
|
||||
track_id=int(row[4]),
|
||||
box=[int(round(row[0])), int(round(row[1])), int(round(row[2])), int(round(row[3]))],
|
||||
score=float(row[5]),
|
||||
det_index=int(row[7]) if len(row) > 7 else None,
|
||||
)
|
||||
)
|
||||
return outputs
|
||||
|
||||
def reset(self) -> None:
|
||||
self._tracker.reset()
|
||||
Loading…
Add table
Add a link
Reference in a new issue