ascii-photobooth/photobooth.py

#!/usr/bin/env python3
#
# ASCII Photobooth live preview - curses TUI front end.
#
# Grabs frames directly via OpenCV/v4l2 (no fswebcam process-spawn per
# frame) and still shells out to the ascii-image-converter binary for the
# actual conversion, and to txt2printr.sh to print a capture. Reads the
# same webcam.conf/ascii.conf/keybinds.conf used by the bash scripts.

import curses
import os
import re
import subprocess
import sys
import time
from datetime import datetime

try:
    import cv2
except ImportError:
    sys.exit(
        "photobooth.py requires OpenCV's Python bindings "
        "(pacman: python-opencv). Install it and retry."
    )

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


def load_conf(name):
    cfg = {}
    path = os.path.join(SCRIPT_DIR, name)
    if not os.path.isfile(path):
        return cfg
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, value = line.partition("=")
            key, value = key.strip(), value.strip()
            if len(value) >= 2 and value[0] == value[-1] and value[0] in "\"'":
                value = value[1:-1]
            cfg[key] = value
    return cfg


def as_bool(value, default=False):
    if not value:
        return default
    return value.strip().lower() in ("1", "true", "yes", "on")


class Config:
    def __init__(self):
        webcam = load_conf("webcam.conf")
        ascii_ = load_conf("ascii.conf")
        keys = load_conf("keybinds.conf")
        printer = load_conf("printer.conf")

        self.device = webcam.get("DEVICE", "/dev/video0")
        w, _, h = webcam.get("RESOLUTION", "640x480").partition("x")
        self.cap_width = int(w or 640)
        self.cap_height = int(h or 480)
        self.jpeg_quality = int(webcam.get("JPEG_QUALITY", "85"))
        self.default_brightness = webcam.get("DEFAULT_BRIGHTNESS", "auto")
        self.default_exposure = webcam.get("DEFAULT_EXPOSURE", "auto")
        self.step = int(webcam.get("STEP", "5"))
        # v4l2-ctl control IDs (snake_case) - a different naming convention
        # than BRIGHTNESS_CONTROL/EXPOSURE_CONTROL, which hold fswebcam's
        # human-readable names for the bash scripts.
        self.brightness_ctrl = webcam.get("BRIGHTNESS_CTRL_ID", "brightness")
        self.exposure_ctrl = webcam.get("EXPOSURE_CTRL_ID", "exposure_absolute")

        self.width = ascii_.get("WIDTH", "").strip()
        self.height = ascii_.get("HEIGHT", "").strip()
        self.complex = as_bool(ascii_.get("COMPLEX"))
        self.char_map = ascii_.get("CHAR_MAP", "").strip()
        self.grayscale = as_bool(ascii_.get("GRAYSCALE"))
        self.negative = as_bool(ascii_.get("NEGATIVE"))
        self.flip_x = as_bool(ascii_.get("FLIP_X"))
        self.flip_y = as_bool(ascii_.get("FLIP_Y"))

        self.printer_a4 = as_bool(printer.get("PRINTER_A4"))
        self.printer_max_lines = int(printer.get("PRINTER_MAX_LINES", "64") or 64)
        self.printer_max_cols = int(printer.get("PRINTER_MAX_COLS", "54") or 54)

        self.key_capture = keys.get("KEY_CAPTURE", " ")
        self.key_brightness_up = keys.get("KEY_BRIGHTNESS_UP", "+")
        self.key_brightness_down = keys.get("KEY_BRIGHTNESS_DOWN", "-")
        self.key_exposure_up = keys.get("KEY_EXPOSURE_UP", ".")
        self.key_exposure_down = keys.get("KEY_EXPOSURE_DOWN", ",")
        self.key_reset = keys.get("KEY_RESET", "r")
        self.key_quit = keys.get("KEY_QUIT", "q")

    def converter_flags(self):
        flags = []
        if self.char_map:
            flags += ["--map", self.char_map]
        elif self.complex:
            flags.append("--complex")
        if self.grayscale:
            flags.append("--grayscale")
        if self.negative:
            flags.append("--negative")
        if self.flip_x:
            flags.append("--flipX")
        if self.flip_y:
            flags.append("--flipY")
        return flags


def convert_to_ascii(jpg_bytes, width, height, flags):
    proc = subprocess.run(
        ["ascii-image-converter", "-", "-d", f"{width},{height}", *flags],
        input=jpg_bytes,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
    )
    return proc.stdout.decode("utf-8", "replace").splitlines()


_ctrl_ranges = {}


def ctrl_range(device, name):
    if name in _ctrl_ranges:
        return _ctrl_ranges[name]
    try:
        out = subprocess.run(
            ["v4l2-ctl", "-d", device, "--list-ctrls"],
            capture_output=True, text=True, timeout=2,
        ).stdout
    except (subprocess.SubprocessError, FileNotFoundError):
        out = ""
    rng = None
    for line in out.splitlines():
        line = line.strip()
        if line.startswith(name + " ") or line.startswith(name + "\t"):
            m = re.search(r"min=(-?\d+) max=(-?\d+)", line)
            if m:
                rng = (int(m.group(1)), int(m.group(2)))
            break
    _ctrl_ranges[name] = rng
    return rng


def apply_ctrl_pct(device, name, pct):
    rng = ctrl_range(device, name)
    if rng is None:
        return
    lo, hi = rng
    val = round(lo + (hi - lo) * pct / 100)
    subprocess.run(
        ["v4l2-ctl", "-d", device, "--set-ctrl", f"{name}={val}"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )


def clamp_pct(value):
    return max(0, min(100, value))


def adjust(current, delta):
    base = 50 if current == "auto" else int(current)
    return clamp_pct(base + delta)


def matches(key, bound):
    return bool(key) and bool(bound) and key.lower() == bound.lower()


def capture_label(cfg):
    return "space" if cfg.key_capture == " " else cfg.key_capture


def box_dims_for(term_rows, term_cols):
    return max(term_rows - 2, 3), max(term_cols, 10)


_CHAR_ASPECT = 0.5  # width:height ratio of a typical monospace glyph cell


def visual_aspect(cols, rows):
    """Physical width:height ratio an ascii grid of cols x rows renders as."""
    return (cols / rows) * _CHAR_ASPECT if rows else 1.0


def crop_to_aspect(frame, aspect):
    """Center-crop a (h, w, ...) frame to the given width:height ratio."""
    h, w = frame.shape[:2]
    if w / h > aspect:
        new_w = max(1, round(h * aspect))
        x0 = (w - new_w) // 2
        return frame[:, x0:x0 + new_w]
    new_h = max(1, round(w / aspect))
    y0 = (h - new_h) // 2
    return frame[y0:y0 + new_h, :]


def prepare_print_frame(frame, cfg, inner_cols, inner_rows):
    """Crop (and rotate, for A4) a raw frame ahead of the print conversion.

    Order matters for intelligibility: crop to roughly the right aspect
    ratio *before* rotating, then rotate the already-correctly-shaped image,
    and only then hand it to ascii-image-converter - rather than converting
    first and mangling already-rendered ascii text, which doesn't account
    for non-square character cells and reads as visibly distorted.
    """
    if cfg.printer_a4:
        cols, rows = cfg.printer_max_cols, cfg.printer_max_lines
        # Cropping happens before the 90-degree rotation below, which swaps
        # width and height, so crop to the inverse of the final aspect ratio.
        frame = crop_to_aspect(frame, 1 / visual_aspect(cols, rows))
        frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
    else:
        cols, rows = inner_cols, inner_rows
        frame = crop_to_aspect(frame, visual_aspect(cols, rows))
    return frame, cols, rows


def main(stdscr):
    cfg = Config()
    conv_flags = cfg.converter_flags()

    curses.curs_set(0)
    stdscr.nodelay(True)

    cap = cv2.VideoCapture(cfg.device, cv2.CAP_V4L2)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, cfg.cap_width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cfg.cap_height)
    if not cap.isOpened():
        raise SystemExit(f"Could not open camera device {cfg.device}")

    brightness, exposure = cfg.default_brightness, cfg.default_exposure
    if brightness != "auto":
        apply_ctrl_pct(cfg.device, cfg.brightness_ctrl, int(brightness))
    if exposure != "auto":
        apply_ctrl_pct(cfg.device, cfg.exposure_ctrl, int(exposure))

    frame_times = []
    message, message_until = "", 0.0
    rows, cols = stdscr.getmaxyx()
    box_dims = box_dims_for(rows, cols)
    box_win = curses.newwin(box_dims[0], box_dims[1], 0, 0)

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                continue

            now = time.time()
            frame_times.append(now)
            frame_times = [t for t in frame_times if now - t < 1.0]
            fps = len(frame_times)

            rows, cols = stdscr.getmaxyx()
            box_rows, box_cols = box_dims_for(rows, cols)
            inner_rows, inner_cols = max(box_rows - 2, 1), max(box_cols - 2, 1)

            if (box_rows, box_cols) != box_dims:
                box_win = curses.newwin(box_rows, box_cols, 0, 0)
                box_dims = (box_rows, box_cols)
                stdscr.erase()

            # Live preview: always upright, cropped to roughly the box's own
            # aspect ratio - independent of PRINTER_A4, whose rotated/resized
            # capture is done separately at capture time (see below).
            preview_frame = crop_to_aspect(frame, visual_aspect(inner_cols, inner_rows))
            _, buf = cv2.imencode(
                ".jpg", preview_frame, [cv2.IMWRITE_JPEG_QUALITY, cfg.jpeg_quality]
            )
            width = cfg.width or str(inner_cols)
            height = cfg.height or str(inner_rows)
            ascii_lines = convert_to_ascii(buf.tobytes(), width, height, conv_flags)

            box_win.erase()
            box_win.box()
            for i, line in enumerate(ascii_lines[:inner_rows]):
                try:
                    box_win.addstr(1 + i, 1, line[:inner_cols])
                except curses.error:
                    pass
            box_win.overwrite(stdscr)

            if now < message_until:
                status = f" {message} "
            else:
                status = f" brightness: {brightness}   exposure: {exposure}   fps: {fps} "
            help_line = (
                f" [{cfg.key_brightness_up}/{cfg.key_brightness_down}] brightness  "
                f"[{cfg.key_exposure_up}/{cfg.key_exposure_down}] exposure  "
                f"[{cfg.key_reset}] reset  [{capture_label(cfg)}] capture  "
                f"[{cfg.key_quit}] quit "
            )
            try:
                stdscr.addstr(box_rows, 0, status[:cols].ljust(cols))
                stdscr.addstr(box_rows + 1, 0, help_line[:cols].ljust(cols - 1))
            except curses.error:
                pass
            stdscr.refresh()

            ch = stdscr.getch()
            if ch == -1:
                continue
            key = chr(ch) if 0 <= ch < 256 else ""

            if matches(key, cfg.key_quit):
                break
            elif matches(key, cfg.key_brightness_up):
                brightness = adjust(brightness, cfg.step)
                apply_ctrl_pct(cfg.device, cfg.brightness_ctrl, brightness)
            elif matches(key, cfg.key_brightness_down):
                brightness = adjust(brightness, -cfg.step)
                apply_ctrl_pct(cfg.device, cfg.brightness_ctrl, brightness)
            elif matches(key, cfg.key_exposure_up):
                exposure = adjust(exposure, cfg.step)
                apply_ctrl_pct(cfg.device, cfg.exposure_ctrl, exposure)
            elif matches(key, cfg.key_exposure_down):
                exposure = adjust(exposure, -cfg.step)
                apply_ctrl_pct(cfg.device, cfg.exposure_ctrl, exposure)
            elif matches(key, cfg.key_reset):
                brightness, exposure = cfg.default_brightness, cfg.default_exposure
                if brightness != "auto":
                    apply_ctrl_pct(cfg.device, cfg.brightness_ctrl, int(brightness))
                if exposure != "auto":
                    apply_ctrl_pct(cfg.device, cfg.exposure_ctrl, int(exposure))
            elif matches(key, cfg.key_capture):
                timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
                photo_path = os.path.join(SCRIPT_DIR, f"photo-{timestamp}.jpg")
                ascii_path = os.path.join(SCRIPT_DIR, f"ascii-img-{timestamp}.txt")

                print_frame, print_cols, print_rows = prepare_print_frame(
                    frame, cfg, inner_cols, inner_rows
                )
                _, print_buf = cv2.imencode(
                    ".jpg", print_frame, [cv2.IMWRITE_JPEG_QUALITY, cfg.jpeg_quality]
                )
                print_jpg_bytes = print_buf.tobytes()
                print_width = cfg.width or str(print_cols)
                print_height = cfg.height or str(print_rows)
                print_ascii_lines = convert_to_ascii(
                    print_jpg_bytes, print_width, print_height, conv_flags
                )

                with open(photo_path, "wb") as f:
                    f.write(print_jpg_bytes)
                with open(ascii_path, "w") as f:
                    f.write("\n".join(print_ascii_lines) + "\n")

                env = os.environ.copy()
                if cfg.printer_a4:
                    # The image was already rotated above; tell txt2printr.sh
                    # not to also apply its own text-transpose rotation.
                    env["ALREADY_ROTATED"] = "true"
                subprocess.run(
                    [os.path.join(SCRIPT_DIR, "txt2printr.sh"), ascii_path], env=env
                )
                message = f"Saved {os.path.basename(photo_path)}"
                message_until = time.time() + 2
    finally:
        cap.release()


if __name__ == "__main__":
    curses.wrapper(main)