Object detection

"""Live object detection from a camera with OpenCV.

usage: python camera.py DETECTOR.slab LICENSE_FILE [--camera 0] [--score 0.4]
                        [--seconds N] [--snapshot out.jpg]

Press q or Esc to quit. Needs OpenCV (`opencv-python`).
"""

import argparse
import time

import cv2

import slab

COCO = ("person bicycle car motorcycle airplane bus train truck boat traffic_light fire_hydrant "
        "stop_sign parking_meter bench bird cat dog horse sheep cow elephant bear zebra giraffe "
        "backpack umbrella handbag tie suitcase frisbee skis snowboard sports_ball kite baseball_bat "
        "baseball_glove skateboard surfboard tennis_racket bottle wine_glass cup fork knife spoon bowl "
        "banana apple sandwich orange broccoli carrot hot_dog pizza donut cake chair couch potted_plant "
        "bed dining_table toilet tv laptop mouse remote keyboard cell_phone microwave oven toaster sink "
        "refrigerator book clock vase scissors teddy_bear hair_drier toothbrush").split()

parser = argparse.ArgumentParser()
parser.add_argument("model")
parser.add_argument("license")
parser.add_argument("--camera", type=int, default=0)
parser.add_argument("--score", type=float, default=0.4)
parser.add_argument("--seconds", type=float, default=0, help="stop after this many seconds (0 = until q)")
parser.add_argument("--snapshot", help="save the last annotated frame to this file")
args = parser.parse_args()

cfg = slab.Config(cache_dir="/tmp/slab_cache", license=open(args.license).read())
detector = slab.Session(cfg, args.model, warmup=True)
labels = detector.labels or (COCO if detector.outputs[-1].shape[1] == 80 else [])

camera = cv2.VideoCapture(args.camera)
if not camera.isOpened():
    raise SystemExit("Cannot open the camera (check camera permission for this terminal)")

inp = slab.Input()
out = slab.Output()
started = time.time()
shown = 0
fps = 0.0
image = None
while True:
    ok, image = camera.read()
    if not ok:
        break
    tick = time.time()

    # OpenCV frames are BGR pixel arrays; Slab copies them into a Frame and
    # prepares them the way the model expects (resize, color order, scale).
    h, w = image.shape[:2]
    frame = slab.Frame.from_pixels(image, w, h, slab.PixelFormat.BGR8)
    inp.clear()
    inp.set_frame(frame)
    detector.run(inp, out)
    boxes = out.boxes(score_threshold=args.score, iou_threshold=0.45, max_boxes=100)

    # Boxes are already in camera-image pixels.
    for b in boxes:
        name = b.label or (labels[b.class_id] if 0 <= b.class_id < len(labels) else str(b.class_id))
        p1, p2 = (int(b.x1), int(b.y1)), (int(b.x2), int(b.y2))
        cv2.rectangle(image, p1, p2, (0, 255, 0), 2)
        cv2.putText(image, f"{name} {b.score:.2f}", (p1[0], max(p1[1] - 6, 12)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        for x, y in b.landmarks:
            cv2.circle(image, (int(x), int(y)), 2, (0, 0, 255), -1)

    fps = 0.9 * fps + 0.1 / max(time.time() - tick, 1e-6) if shown else 1.0 / max(time.time() - tick, 1e-6)
    cv2.putText(image, f"{out.backend_used.name} {out.latency_ms:.1f} ms  {fps:.0f} fps  {len(boxes)} objects",
                (10, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
    cv2.imshow("Slab", image)
    shown += 1

    key = cv2.waitKey(1) & 0xFF
    if key in (ord("q"), 27) or (args.seconds and time.time() - started > args.seconds):
        break

if args.snapshot and image is not None:
    cv2.imwrite(args.snapshot, image)
print(f"{shown} frames in {time.time() - started:.1f} s; last run {out.latency_ms:.1f} ms on {out.backend_used.name}")
camera.release()
cv2.destroyAllWindows()