motionlab/main.py

"""motionlab - a tiny TouchDesigner-style playground: live camera in, motion-reactive graphics out."""
import argparse
import os
import sys
import time

import cv2
import numpy as np

from effects import Context, default_layers, hue_bgr
from panel import Panel
from shaders import ShaderBank

SIZE = (960, 540)
VCAM_SIZE = (1280, 720)
WINDOW = "motionlab"


class Synthetic:
    """Drifting blobs, so the app (and --headless) works without a camera."""

    def __init__(self):
        self.i = 0

    def read(self):
        w, h = SIZE
        t = self.i / 30
        self.i += 1
        img = np.full((h, w, 3), (30, 20, 20), np.uint8)
        for k in range(5):
            x = int(w * (0.5 + 0.4 * np.sin(t * (0.7 + 0.2 * k) + k)))
            y = int(h * (0.5 + 0.4 * np.cos(t * (0.5 + 0.3 * k) + 2 * k)))
            cv2.circle(img, (x, y), 40 + 12 * k, hue_bgr(k / 5) * 255, -1)
        return True, cv2.GaussianBlur(img, (0, 0), 4)

    def release(self):
        pass


def open_source(args):
    if args.demo or args.headless:
        return Synthetic()
    backend = cv2.CAP_AVFOUNDATION if sys.platform == "darwin" else cv2.CAP_ANY
    cap = cv2.VideoCapture(args.camera, backend)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    # without camera permission macOS can hand back all-black frames instead of failing
    for _ in range(10):
        ok, frame = cap.read()
        if ok and frame.any():
            return cap
    cap.release()
    print("Could not read from the camera - falling back to the synthetic demo source.\n"
          "On macOS: allow your terminal under System Settings > Privacy & Security > Camera, then rerun.",
          file=sys.stderr)
    return Synthetic()


def open_virtualcam():
    """Virtual camera other apps (Meet, Zoom, ...) can pick as their webcam."""
    import pyvirtualcam
    try:
        cam = pyvirtualcam.Camera(*VCAM_SIZE, 30, fmt=pyvirtualcam.PixelFormat.BGR)
    except RuntimeError as e:
        sys.exit(f"Could not start the virtual camera: {e}")
    print(f"virtual camera on: pick '{cam.device}' as the camera in your call app")
    return cam


def fit(frame):
    """Centre-crop to the output aspect ratio, then resize."""
    w, h = SIZE
    fh, fw = frame.shape[:2]
    cw = min(fw, fh * w // h)
    ch = min(fh, fw * h // w)
    x, y = (fw - cw) // 2, (fh - ch) // 2
    return cv2.resize(frame[y:y + ch, x:x + cw], SIZE, interpolation=cv2.INTER_AREA)


def capture_path(ext):
    os.makedirs("captures", exist_ok=True)
    return os.path.join("captures", time.strftime(f"motionlab-%Y%m%d-%H%M%S.{ext}"))


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--camera", type=int, default=0, help="camera index (default 0)")
    ap.add_argument("--demo", action="store_true", help="use the synthetic source instead of the camera")
    ap.add_argument("--virtualcam", action="store_true", help="also send the output to a virtual camera for video calls")
    ap.add_argument("--headless", type=int, metavar="N", help="render N synthetic frames without a window, then exit")
    ap.add_argument("--out", default="headless.png", help="where --headless writes its last frame")
    ap.add_argument("--effects", default="effects", help="folder of .glsl shader effects (default ./effects)")
    ap.add_argument("--only", metavar="NAMES", help="enable only these layers, comma separated (e.g. --only liquid,halftone)")
    args = ap.parse_args()

    w, h = SIZE
    src = open_source(args)
    layers = default_layers()
    # prefer ./effects so edits are picked up live; fall back to the copy shipped next to this file
    effects_dir = args.effects if os.path.isdir(args.effects) else os.path.join(os.path.dirname(__file__), "effects")
    bank = ShaderBank(effects_dir, layers)
    if args.only:
        wanted = {name.strip().lower() for name in args.only.split(",")}
        unknown = wanted - {layer.name.lower() for layer in layers}
        if unknown:
            sys.exit(f"unknown layer(s): {', '.join(sorted(unknown))}. Available: {', '.join(l.name for l in layers)}")
        for layer in layers:
            layer.enabled = layer.name.lower() in wanted
    if args.headless and bank.errors:
        sys.exit(f"shader compile failed: {', '.join(bank.errors)}")
    ctx = Context(SIZE)
    panel = Panel(layers)
    canvas = np.empty((h, w + Panel.WIDTH, 3), np.uint8)
    show_panel, writer, fps = True, None, 30.0
    vcam = open_virtualcam() if args.virtualcam else None

    if not args.headless:
        cv2.namedWindow(WINDOW, cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
        cv2.resizeWindow(WINDOW, w + Panel.WIDTH, h)
        cv2.setMouseCallback(WINDOW, panel.on_mouse)

    start = last = time.perf_counter()
    frames = 0
    while True:
        ok, frame = src.read()
        if not ok:
            break
        now = time.perf_counter()
        dt = 1 / 30 if args.headless else max(now - last, 1e-3)
        last = now
        fps += (1 / dt - fps) * 0.1

        frame = cv2.flip(fit(frame), 1)  # mirror, so it behaves like a mirror
        if frames % 15 == 0:
            bank.refresh()  # pick up new and edited shader files
        ctx.update(frame, dt)
        out = frame
        for layer in layers:
            if layer.enabled:
                out = layer.apply(out, ctx)

        if writer:
            writer.write(out)
        if vcam:
            # un-mirror: the preview is a mirror, but the people on the call should see the real thing
            vcam.send(cv2.resize(cv2.flip(out, 1), VCAM_SIZE))
        view = out
        if show_panel:
            canvas[:, :w] = out
            panel.draw(canvas, w, f"{fps:3.0f} fps" + ("  REC" if writer else "") + ("  CAM" if vcam else ""))
            view = canvas
        else:
            panel.hits = []

        frames += 1
        if args.headless:
            if frames >= args.headless:
                cv2.imwrite(args.out, view)
                print(f"{frames} frames, {(time.perf_counter() - start) / frames * 1000:.1f} ms/frame -> {args.out}")
                break
            continue

        cv2.imshow(WINDOW, view)
        key = cv2.waitKey(1) & 0xFF
        if key in (ord("q"), 27) or cv2.getWindowProperty(WINDOW, cv2.WND_PROP_VISIBLE) < 1:
            break
        elif key == ord("h"):
            show_panel = not show_panel
        elif key == ord("s"):
            path = capture_path("png")
            cv2.imwrite(path, out)
            print("saved", path)
        elif key == ord("r"):
            if writer:
                writer.release()
                writer = None
            else:
                path = capture_path("mp4")
                writer = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), round(fps), SIZE)
                print("recording", path)

    if writer:
        writer.release()
    if vcam:
        vcam.close()
    src.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()