Blame for server/reading.py

"""Read one picture the way QR Scan Difficulty read every code, with the five of its eight scanners that run on Linux.

As in the measurement:
  * the picture is read in gray and shrunk, never enlarged, so its long side is at most each frame size (1920 down to
    224 pixels), with OpenCV's INTER_AREA; sizes that come out the same are read once;
  * each frame is staged as a gray PNG; zxing-cpp, WeChat's scanner and OpenCV's two detectors read it in a worker
    process (one OpenCV thread, library defaults, QR codes only), and ZXing reads it in a Java process the way
    Android's Settings app calls it;
  * a decode that runs past TIMEOUT_S counts as a miss: the worker is replaced and the rest is read without that call.
No measured picture had transparency; here a transparent part counts as white, as on a printed page.

A scanner reads a frame when one of the contents it returns is the picture's content. The measurement knew each code's
content; here it is the content the most scanners return (ties: the most frames, then text over raw bytes, then the
shortest). Raw bytes and the text a scanner makes of them (ISO-8859-1, QR's default, or Shift JIS) are one content.
"""

from __future__ import annotations

import base64
import hashlib
import itertools
import json
import os
import queue
import subprocess
import tempfile
import threading
import time
from multiprocessing import get_context
from pathlib import Path

os.environ.setdefault("OPENCV_IO_MAX_IMAGE_PIXELS", "50000000")   # OpenCV refuses bigger pictures before decoding them
VENDOR = Path(os.environ.get("QSD_VENDOR", Path(__file__).resolve().parent.parent / "vendor"))
SELF_TEST = Path(__file__).with_name("selftest.png")   # a clean code, which all five read at every size
FRAMES = [1920, 1280, 896, 640, 448, 320, 224]
POOL = ["zxing-cpp", "wechat", "opencv", "opencv-aruco"]   # read in this order, as the measurement's pool read them
SCANNERS = [*POOL, "zxing-java"]
TIMEOUT_S = 20          # one decode, as in the measurement
STAGE_TIMEOUT_S = 30    # opening and shrinking one picture
DEADLINE_S = 40         # one picture in all, so that a request ends within a minute; decodes not done by then count
                        # as misses and the result says "partial"


class PictureError(Exception):
    """A picture this server cannot open."""


class _TooBig(Exception):
    """More pixels than OPENCV_IO_MAX_IMAGE_PIXELS."""


def frame_sizes(h: int, w: int) -> dict[tuple[int, int], list[int]]:
    """{(w, h): [frame sizes]} for the frames, never enlarging (the measurement's rule, rounding included)."""
    sizes: dict[tuple[int, int], list[int]] = {}
    for f in FRAMES:
        k = min(1.0, f / max(h, w))
        sizes.setdefault((max(1, round(w * k)), max(1, round(h * k))), []).append(f)
    return sizes


# ---- the worker process: zxing-cpp, WeChat's scanner, OpenCV's two detectors --------------------------------------

def _decoders(wechat_dir: str):
    import cv2
    import zxingcpp

    cv2.setNumThreads(1)
    d = Path(wechat_dir)
    wechat = cv2.wechat_qrcode.WeChatQRCode(str(d / "detect.prototxt"), str(d / "detect.caffemodel"),
                                            str(d / "sr.prototxt"), str(d / "sr.caffemodel"))
    qr_format = zxingcpp.BarcodeFormat.QRCode   # in zxing-cpp 3 this also returns Micro QR and rMQR, dropped below
    dropped = (zxingcpp.BarcodeFormat.MicroQRCode, zxingcpp.BarcodeFormat.RMQRCode)

    def zxing_cpp(bgr):
        out = []
        for r in zxingcpp.read_barcodes(cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY), formats=qr_format):
            if r.format in dropped or not getattr(r, "valid", True):
                continue
            out.append(bytes(r.bytes) if r.bytes is not None else (r.text or "").encode("utf-8"))
        return out

    def wechat_qr(bgr):
        try:
            texts, _ = wechat.detectAndDecode(bgr)
        except UnicodeDecodeError as e:   # binary content: OpenCV's binding fails to make text of it
            return [bytes(e.object)]
        return [t.encode("utf-8") for t in texts if t]

    def single(detector):   # detectAndDecode, the single-code call apps make
        def read(bgr):
            text, _, _ = detector.detectAndDecode(cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY))
            return [text.encode("utf-8")] if text else []
        return read

    return cv2, {"zxing-cpp": zxing_cpp, "wechat": wechat_qr,
                 "opencv": single(cv2.QRCodeDetector()), "opencv-aruco": single(cv2.QRCodeDetectorAruco())}


def _open(cv2, path: str):
    """The picture in gray as the measurement opened it, any transparent part on white; None if it can't be opened."""
    import numpy as np

    try:
        gray = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        if gray is None or not path.endswith((".png", ".webp")):
            return gray
        full = cv2.imread(path, cv2.IMREAD_UNCHANGED)   # gray (or colour) with alpha comes back as 4 channels
    except cv2.error as e:   # past OPENCV_IO_MAX_IMAGE_PIXELS OpenCV raises, rather than return None
        if "MAX_IMAGE_PIXELS" in str(e):
            raise _TooBig from e
        return None
    if full is None or full.ndim != 3 or full.shape[2] != 4:
        return gray
    top = np.iinfo(full.dtype).max
    if full[..., 3].min() == top:   # no pixel is transparent: the measured path, unchanged
        return gray
    a = full[..., 3].astype(np.float32) / top
    g = cv2.cvtColor(full[..., :3], cv2.COLOR_BGR2GRAY).astype(np.float32) / top
    return np.rint(255 * (g * a + (1 - a))).astype(np.uint8)


def _child(conn, wechat_dir: str) -> None:
    try:
        cv2, decoders = _decoders(wechat_dir)
    except Exception as e:  # noqa: BLE001 - reported to the parent, which refuses to start
        conn.send(("broken", repr(e)))
        return
    conn.send(("ready",))
    while True:
        try:
            msg = conn.recv()
        except EOFError:
            return
        if msg[0] == "stage":
            _, path, outdir = msg
            try:
                gray = _open(cv2, path)
                if gray is None:
                    conn.send(("unreadable",))
                    continue
                frames = []
                for k, ((w, h), labels) in enumerate(frame_sizes(*gray.shape).items()):
                    g = gray if (h, w) == gray.shape else cv2.resize(gray, (w, h), interpolation=cv2.INTER_AREA)
                    p = os.path.join(outdir, f"{k}.png")
                    cv2.imwrite(p, g, [cv2.IMWRITE_PNG_COMPRESSION, 1])
                    frames.append((w, h, labels, p))
                conn.send(("staged", gray.shape[1], gray.shape[0], frames))
            except _TooBig:
                conn.send(("too big",))
            except Exception:  # noqa: BLE001 - a picture that can't be staged (out of memory, say) is refused
                conn.send(("unreadable",))
        elif msg[0] == "decode":
            _, frames, skip = msg
            for k, (_, _, _, p) in enumerate(frames):
                todo = [n for n in POOL if (k, n) not in skip]
                if not todo:
                    continue
                bgr = cv2.cvtColor(cv2.imread(p, cv2.IMREAD_GRAYSCALE), cv2.COLOR_GRAY2BGR)
                for n in todo:
                    conn.send(("start", k, n))
                    try:
                        payloads = decoders[n](bgr)
                    except Exception:  # noqa: BLE001 - a decoder that fails reads nothing, as in the measurement
                        payloads = []
                    conn.send(("done", k, n, payloads))
            conn.send(("end",))


class _Worker:
    """One decoder worker process. A stuck one is killed; a new one starts when the next message is sent."""

    def __init__(self, ctx):
        self._ctx = ctx
        self._spawn()

    def _spawn(self) -> None:
        self.conn, child = self._ctx.Pipe()
        self.proc = self._ctx.Process(target=_child, args=(child, str(VENDOR / "wechat")), daemon=True)
        self.proc.start()
        child.close()
        msg = self._recv(180)
        if msg is None or msg[0] != "ready":
            raise RuntimeError(f"decoder worker did not start: {msg}")

    def close(self) -> None:
        self.proc.kill()
        self.proc.join(5)
        self.conn.close()

    def _send(self, msg) -> None:
        """Send to a live worker: one that was killed, or died while idle (killed for memory, say), is replaced."""
        if self.conn.closed or not self.proc.is_alive():
            self.close()
            self._spawn()
        try:
            self.conn.send(msg)
        except OSError:   # it died just now
            self.close()
            self._spawn()
            self.conn.send(msg)

    def _recv(self, timeout: float):
        try:
            if not self.conn.poll(timeout):
                return None
            return self.conn.recv()
        except (EOFError, OSError):
            return None

    def stage(self, path: str, outdir: str):
        t0 = time.monotonic()
        self._send(("stage", path, outdir))
        msg = self._recv(STAGE_TIMEOUT_S)
        if msg is None:   # too slow, or the worker crashed on this picture
            slow = time.monotonic() - t0 >= STAGE_TIMEOUT_S
            self.close()
            raise PictureError("This picture took too long to open." if slow else "This picture could not be opened.")
        if msg[0] == "too big":
            raise PictureError("This picture is over 50 megapixels.")
        if msg[0] == "unreadable":
            raise PictureError("This picture could not be opened.")
        return msg[1], msg[2], msg[3]

    def decode(self, frames, deadline: float):
        """{(frame index, scanner): payloads}, and whether the deadline (or crashes) cut the reading short."""
        got: dict[tuple[int, str], list[bytes]] = {}
        missed: set[tuple[int, str]] = set()
        crashes = 0
        while True:
            self._send(("decode", frames, missed | set(got)))
            current = None
            while True:
                msg = self._recv(max(0.0, min(TIMEOUT_S, deadline - time.monotonic())))
                if msg is None:
                    break
                if msg[0] == "start":
                    current = (msg[1], msg[2])
                elif msg[0] == "done":
                    got[(msg[1], msg[2])] = msg[3]
                    current = None
                elif msg[0] == "end":
                    return got, False
            # silence: a decode ran past TIMEOUT_S (a miss), the worker died, or the deadline passed
            self.close()
            if time.monotonic() >= deadline:
                return got, True
            if current is not None:
                missed.add(current)
            else:
                crashes += 1
                if crashes > 1:
                    return got, True


# ---- the Java process: ZXing as Android's Settings app calls it ---------------------------------------------------

class _Java:
    def __init__(self):
        cp = os.pathsep.join([str(VENDOR / "zxing" / "core-3.5.2.jar"), str(VENDOR / "zxing" / "classes")])
        self.cmd = ["java", "-Djava.awt.headless=true", "-Xmx512m", "-cp", cp, "ZxingServer"]
        self.lock = threading.Lock()
        self.ids = itertools.count()
        self._spawn()

    def _spawn(self) -> None:
        self.proc = subprocess.Popen(self.cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                     stderr=subprocess.DEVNULL, text=True, encoding="utf-8", bufsize=1)
        self.lines: queue.Queue = queue.Queue()
        threading.Thread(target=self._pump, args=(self.proc.stdout, self.lines), daemon=True).start()

    @staticmethod
    def _pump(stream, lines: queue.Queue) -> None:
        for line in stream:
            lines.put(line)
        lines.put(None)

    def close(self) -> None:
        self.proc.kill()
        self.proc.wait(5)
        try:
            self.proc.stdin.close()
        except OSError:   # its pipe broke with it
            pass

    def read(self, paths: list[str], deadline: float) -> list[list[bytes] | None]:
        """Per frame, ZXing's contents (empty: no read, a timeout included); None where the deadline left it unread."""
        out: list[list[bytes] | None] = [None] * len(paths)
        deaths = 0
        for k, p in enumerate(paths):
            # one frame at a time, so two pictures read at once take turns; never waiting past this one's deadline
            if not self.lock.acquire(timeout=max(0.0, deadline - time.monotonic())):
                break
            try:
                left = deadline - time.monotonic()
                if left <= 0:
                    break
                if self.proc.poll() is not None:   # it died since its last answer (while idle): not a miss
                    self.close()
                    self._spawn()
                rid = next(self.ids)
                obj, died = None, False
                try:
                    self.proc.stdin.write(f"{rid}\t{p}\n")
                    self.proc.stdin.flush()
                    line = self.lines.get(timeout=min(TIMEOUT_S, left))
                    died = line is None   # its output ended: it exited without answering
                    obj = json.loads(line) if line else None
                except OSError:   # its pipe broke: it exited
                    died = True
                except (ValueError, queue.Empty):
                    pass
                if not isinstance(obj, dict) or obj.get("id") != rid:   # a timeout, a crash or a stray line: a miss
                    self.close()
                    self._spawn()
                    deaths += died
                    if deaths > 1:   # it keeps dying: this frame and the rest are unread, not misses
                        break
                    out[k] = [] if time.monotonic() < deadline else None
                    continue
                out[k] = [base64.b64decode(c) for c in obj.get("contents", [])]
            finally:
                self.lock.release()
        return out


# ---- one picture ----------------------------------------------------------------------------------------------------

def _texts(p: bytes) -> set[bytes]:
    """In UTF-8, the texts a scanner may make of raw bytes that aren't UTF-8: ISO-8859-1 (QR byte mode's default)
    and, where they decode, Shift JIS (QR kanji mode)."""
    try:
        p.decode("utf-8")
        return set()
    except UnicodeDecodeError:
        pass
    out = {p.decode("latin-1").encode("utf-8")}
    try:
        out.add(p.decode("shift_jis").encode("utf-8"))
    except UnicodeDecodeError:
        pass
    return out


def _same(p: bytes, q: bytes) -> bool:
    """One content: the same bytes, or raw bytes and a text made of them."""
    return p == q or p in _texts(q) or q in _texts(p)


def _is_text(p: bytes) -> bool:
    try:
        p.decode("utf-8")
        return True
    except UnicodeDecodeError:
        return False


def _consensus(contents: dict[tuple[str, int], list[bytes]]) -> bytes | None:
    forms = {p for ps in contents.values() for p in ps}
    if not forms:
        return None

    def votes(c: bytes) -> tuple[int, int]:
        hits = [(s, k) for (s, k), ps in contents.items() if any(_same(p, c) for p in ps)]
        return len({s for s, _ in hits}), len(hits)

    return max(forms, key=lambda c: (*votes(c), _is_text(c), -len(c), c))


def describe(content: bytes | None) -> dict | None:
    """The content as the page may show it: printable text (UTF-8, or else ISO-8859-1), otherwise only its length."""
    if content is None:
        return None
    try:
        text = content.decode("utf-8")
    except UnicodeDecodeError:
        text = content.decode("latin-1")   # QR byte mode's default; binary content shows control characters
    if all(c.isprintable() or c in "\r\n\t" for c in text):
        return {"text": text[:300], "cut": len(text) > 300, "bytes": len(content)}
    return {"text": None, "cut": False, "bytes": len(content)}


class Reader:
    """Reads pictures, `slots` at a time (one decoder worker each); the ZXing process is shared."""

    def __init__(self, slots: int = 2):
        ctx = get_context("spawn")
        self.java = _Java()
        self.workers = [_Worker(ctx) for _ in range(slots)]
        self.free: queue.Queue = queue.Queue()
        for w in self.workers:
            self.free.put(w)
        self._self_test()

    def _self_test(self) -> None:
        """Refuse to start, rather than place every picture wrongly, unless all five read a clean code at every size."""
        r = self.read(str(SELF_TEST))
        missed = [n for n in SCANNERS if not all(r["reads"][n])]
        if missed or r["partial"]:
            self.close()
            raise RuntimeError(f"self-test failed: {', '.join(missed) or 'the reading ran out of time'}")

    def close(self) -> None:
        for w in self.workers:
            w.close()
        self.java.close()

    def read(self, path: str, truth: set[str] | None = None) -> dict:
        """Read one picture file. `truth` (sha256 hex digests of the content) replaces the consensus, as the
        measurement scored known codes; the server never passes it."""
        t0 = time.monotonic()
        deadline = t0 + DEADLINE_S
        worker = self.free.get()
        try:
            with tempfile.TemporaryDirectory(prefix="qsd-") as tmp:
                w, h, frames = worker.stage(path, tmp)
                java: dict = {"r": [None] * len(frames)}

                def zxing_java():
                    java["r"] = self.java.read([f[3] for f in frames], deadline)

                t = threading.Thread(target=zxing_java, daemon=True)
                t.start()
                try:
                    pool, partial = worker.decode(frames, deadline)
                finally:
                    t.join()
        finally:
            self.free.put(worker)
        contents = {(n, k): p for (k, n), p in pool.items()}
        for k, p in enumerate(java["r"]):
            contents[("zxing-java", k)] = p or []
            partial = partial or p is None
        if truth is None:
            content = _consensus(contents)
            def hit(ps): return content is not None and any(_same(p, content) for p in ps)  # noqa: E306
        else:
            content = None
            def hit(ps): return any(hashlib.sha256(p).hexdigest() in truth for p in ps)  # noqa: E306
        reads = {n: [0] * len(FRAMES) for n in SCANNERS}
        sizes = [None] * len(FRAMES)
        for k, (fw, fh, labels, _) in enumerate(frames):
            for f in labels:
                fi = FRAMES.index(f)
                sizes[fi] = [fw, fh]
                for n in SCANNERS:
                    reads[n][fi] = int(hit(contents.get((n, k), [])))
        return {"size": [w, h], "frames": sizes, "reads": reads, "content": describe(content), "partial": partial,
                "ms": round(1000 * (time.monotonic() - t0))}