diff --git a/README.md b/README.md index 7932464..7311f91 100644 --- a/README.md +++ b/README.md @@ -22,8 +22,9 @@ Linux driver tools for the Hama myScan (00052343) scanner mouse, a Dacuda design - `tools/app.py` — desktop app (Qt/PySide6; Linux, Windows, macOS): live camera and page view while scanning (Start/Stop or the scan button, F5), then OCR of the finished page with RapidOCR (`myscan/ocr.py`; multilingual PP-OCRv6 models bundled in the wheel, works offline). Recognised lines are outlined on the page by - confidence; click one to select its text, and the cursor in the text highlights its line. The text is editable - and can be copied or saved as .txt. `python tools/app.py [image]` also opens and reads existing images (Open, or + confidence, and the text is selected right on the page like in a PDF viewer: drag (in reading order across + lines), double-click a word, triple-click a line, shift-click to extend, Ctrl+A; Ctrl+C or right-click copies. + Middle-drag or Space+drag pans, the wheel zooms. Copy all and Save text (.txt) take the whole page. `python tools/app.py [image]` also opens and reads existing images (Open, or drag and drop), which is all it does on machines without the mouse - `myscan/scanner.py` — the scan loop shared by `scan.py` and the app (finds the mouse's `/dev/videoN` itself) - `myscan/live.py`, `myscan/webview.py` — the live preview and its browser viewer (standard library only) diff --git a/myscan/app.py b/myscan/app.py index d77bb41..e81c6f3 100644 --- a/myscan/app.py +++ b/myscan/app.py @@ -2,26 +2,28 @@ Left, the camera and the page as it is being stitched (the same preview as the browser viewer); after a scan the finished page is loaded and read with RapidOCR (myscan/ocr.py). -Recognised lines are outlined on the page, coloured by confidence; clicking one selects -its text on the right, and moving the cursor in the text highlights its line on the -page. The text can be edited, copied and saved. Images can also be opened or dropped -onto the window, so OCR works on any platform; scanning needs Linux and the mouse. +Recognised lines are outlined on the page, coloured by confidence, and their text can be +selected right on the page (drag, double-click a word, triple-click a line) and copied; +all of it can be copied or saved. Images can also be opened or dropped onto the window, +so OCR works on any platform; scanning needs Linux and the mouse. Run: python tools/app.py [image] """ import os +import signal import sys import threading import time +import traceback import cv2 import numpy as np -from PySide6.QtCore import QObject, QPointF, QRectF, Qt, QThread, Signal -from PySide6.QtGui import (QAction, QBrush, QColor, QGuiApplication, QImage, QKeySequence, QPainter, QPen, - QPixmap, QPolygonF, QTextCursor) -from PySide6.QtWidgets import (QApplication, QFileDialog, QGraphicsPixmapItem, QGraphicsPolygonItem, - QGraphicsScene, QGraphicsView, QLabel, QMainWindow, QMessageBox, QPlainTextEdit, - QSizePolicy, QSplitter, QToolBar, QVBoxLayout, QWidget) +from PySide6.QtCore import QObject, QPointF, QRectF, Qt, QThread, QTimer, Signal +from PySide6.QtGui import (QAction, QBrush, QClipboard, QColor, QGuiApplication, QImage, QKeySequence, QPainter, + QPainterPath, QPen, QPixmap, QPolygonF) +from PySide6.QtWidgets import (QApplication, QFileDialog, QGraphicsPathItem, QGraphicsPixmapItem, + QGraphicsPolygonItem, QGraphicsScene, QGraphicsView, QLabel, QMainWindow, QMenu, + QMessageBox, QSizePolicy, QSplitter, QToolBar, QVBoxLayout, QWidget) from .ocr import Ocr, layout @@ -48,6 +50,13 @@ def imwrite(path, img): buf.tofile(path) +def log_error(msg): + """Print an error, with the traceback of the exception being handled if any.""" + print(f"myScan: {msg}", file=sys.stderr, flush=True) + if sys.exc_info()[0] is not None: + traceback.print_exc() + + def scanner_device(): """The myScan's capture device, or None (always None off Linux).""" if not sys.platform.startswith("linux"): @@ -108,12 +117,14 @@ class ScanThread(QThread): self.bridge.set_status(state="starting", message="reading calibration") scanner = Scanner(self.opts, self.bridge) except Exception as e: + log_error(f"cannot open the scanner: {e}") self.failed.emit(f"cannot open the scanner: {e}") return try: while scanner.wait_press(self.quit): scanner.scan() except Exception as e: + log_error(f"scan failed: {e}") self.failed.emit(f"scan failed: {e}") finally: scanner.close() @@ -137,48 +148,50 @@ class OcrThread(QThread): try: self.done.emit(self.ocr(self.img), time.time() - t) except Exception as e: + log_error(f"OCR failed: {e}") self.failed.emit(str(e)) class LineItem(QGraphicsPolygonItem): - def __init__(self, index, box, score, on_click): - super().__init__(QPolygonF([QPointF(*p) for p in box])) - self.index, self.on_click = index, on_click - c = QColor(46, 160, 67) if score >= 0.9 else QColor(219, 140, 0) if score >= 0.7 else QColor(214, 40, 40) - self.color = c - self.setAcceptHoverEvents(True) - self.setCursor(Qt.PointingHandCursor) - self.setToolTip(f"{score:.0%}") - self.set_active(False) + """Outline of a recognised line, coloured by confidence (mouse events go to the view).""" - def set_active(self, on): - pen = QPen(QColor(30, 110, 230) if on else self.color, 3 if on else 1.5) + def __init__(self, box, score): + super().__init__(QPolygonF([QPointF(*p) for p in box])) + c = QColor(46, 160, 67) if score >= 0.9 else QColor(219, 140, 0) if score >= 0.7 else QColor(214, 40, 40) + pen = QPen(c, 1.5) pen.setCosmetic(True) self.setPen(pen) - fill = QColor(30, 110, 230, 70) if on else QColor(self.color.red(), self.color.green(), self.color.blue(), 25) - self.setBrush(QBrush(fill)) - - def mousePressEvent(self, ev): - self.on_click(self.index) + self.setBrush(QBrush(QColor(c.red(), c.green(), c.blue(), 25))) + self.setAcceptedMouseButtons(Qt.NoButton) + self.setToolTip(f"{score:.0%}") class PageView(QGraphicsView): - """Zoomable, pannable page with the OCR line boxes on top. Keeps the page fitted to - the window until the user zooms.""" - lineClicked = Signal(int) + """Zoomable page whose recognised text can be selected like in a PDF viewer: drag to + select in reading order, double-click for a word, triple-click for a line, shift-click + to extend. Middle-drag (or Space + drag, or any drag before there is text) pans. Keeps + the page fitted to the window until the user zooms.""" + selectionChanged = Signal(str) def __init__(self): super().__init__() self.setScene(QGraphicsScene(self)) self.setRenderHints(QPainter.Antialiasing | QPainter.SmoothPixmapTransform) - self.setDragMode(QGraphicsView.ScrollHandDrag) self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse) self.setBackgroundBrush(QColor(30, 32, 36)) + self.setMouseTracking(True) self.pixmap = QGraphicsPixmapItem() self.pixmap.setTransformationMode(Qt.SmoothTransformation) self.scene().addItem(self.pixmap) - self.items_, self.active, self.auto_fit = [], None, True - self.boxes_visible = True + self.sel_item = QGraphicsPathItem() + self.sel_item.setPen(QPen(Qt.NoPen)) + self.sel_item.setBrush(QColor(30, 110, 230, 90)) + self.sel_item.setZValue(2) + self.scene().addItem(self.sel_item) + self.items_, self.auto_fit, self.boxes_visible = [], True, True + self.lines, self.spans, self.text = [], [], "" + self.sel, self.anchor, self.unit = (0, 0), None, None # unit: drag extends by "word"/"line" + self.pan_from, self.space, self.last_double = None, False, None def set_image(self, qimg, keep_view=False): self.pixmap.setPixmap(QPixmap.fromImage(qimg)) @@ -188,28 +201,150 @@ class PageView(QGraphicsView): if self.auto_fit: self.fit() - def set_lines(self, lines): + def set_lines(self, lines, text="", spans=()): for it in self.items_: self.scene().removeItem(it) - self.items_ = [LineItem(i, ln.box, ln.score, self.lineClicked.emit) for i, ln in enumerate(lines)] + self.items_ = [LineItem(ln.box, ln.score) for ln in lines] for it in self.items_: it.setVisible(self.boxes_visible) self.scene().addItem(it) - self.active = None + self.lines, self.text, self.spans = lines, text, list(spans) + self.sel, self.anchor = None, None + self.select(0, 0) def show_boxes(self, on): self.boxes_visible = on for it in self.items_: it.setVisible(on) - def highlight(self, i): - if self.active is not None and self.active < len(self.items_): - self.items_[self.active].set_active(False) - self.active = i - if i is not None and i < len(self.items_): - self.items_[i].set_active(True) - if not self.auto_fit: - self.ensureVisible(self.items_[i], 40, 40) + # ---- selection ---- + + def hit(self, scene_pos): + """(line index, distance in pixels, text position) of the line nearest a scene point.""" + p = (scene_pos.x(), scene_pos.y()) + return min(((i, *ln.locate(p)) for i, ln in enumerate(self.lines)), key=lambda h: h[1]) + + def position(self, scene_pos, unit=None): + """Range in the text at a scene point: empty, or the word/line around it.""" + i, _, k = self.hit(scene_pos) + s = self.spans[i][0] + if unit == "word": + a, b = self.lines[i].word(k) + elif unit == "line": + a, b = 0, len(self.lines[i].text) + else: + a = b = k + return s + a, s + b + + def select(self, a, b): + a, b = min(a, b), max(a, b) + if (a, b) == self.sel: + return + self.sel = (a, b) + path = QPainterPath() + for ln, (s, e) in zip(self.lines, self.spans): + lo, hi = max(a, s), min(b, e) + if lo < hi: + path.addPolygon(QPolygonF([QPointF(*p) for p in ln.quad(lo - s, hi - s)])) + path.closeSubpath() + self.sel_item.setPath(path) + self.selectionChanged.emit(self.selected_text()) + + def select_all(self): + self.select(0, len(self.text)) + + def selected_text(self): + a, b = self.sel + return self.text[a:b] + + def drag_to(self, scene_pos): + a, b = self.position(scene_pos, self.unit) + a0, b0 = self.anchor + self.select(min(a, a0), max(b, b0)) + + # ---- mouse ---- + + def mousePressEvent(self, ev): + selecting = ev.button() == Qt.LeftButton and self.lines and not self.space + if not selecting: + if ev.button() in (Qt.LeftButton, Qt.MiddleButton): + self.pan_from = ev.position() + self.viewport().setCursor(Qt.ClosedHandCursor) + return super().mousePressEvent(ev) + pos = self.mapToScene(ev.position().toPoint()) + triple = (self.last_double is not None and time.monotonic() - self.last_double[0] + < QApplication.styleHints().mouseDoubleClickInterval() / 1000 + and (ev.position() - self.last_double[1]).manhattanLength() < 6) + self.unit = "line" if triple else None + r = self.position(pos, self.unit) + if ev.modifiers() & Qt.ShiftModifier and self.anchor is not None: + self.drag_to(pos) + else: + self.anchor = r + self.select(*r) + + def mouseDoubleClickEvent(self, ev): + if ev.button() != Qt.LeftButton or not self.lines or self.space: + return super().mouseDoubleClickEvent(ev) + self.last_double = (time.monotonic(), ev.position()) + self.unit = "word" + self.anchor = self.position(self.mapToScene(ev.position().toPoint()), "word") + self.select(*self.anchor) + + def mouseMoveEvent(self, ev): + if self.pan_from is not None: + d = ev.position() - self.pan_from + self.pan_from = ev.position() + self.horizontalScrollBar().setValue(self.horizontalScrollBar().value() - round(d.x())) + self.verticalScrollBar().setValue(self.verticalScrollBar().value() - round(d.y())) + return + pos = self.mapToScene(ev.position().toPoint()) + if ev.buttons() & Qt.LeftButton and self.anchor is not None and self.lines: + self.drag_to(pos) + self.ensureVisible(QRectF(pos, pos), 0, 0) # scroll when dragging past the edge + elif self.lines and not self.space: + i, dist, _ = self.hit(pos) + near = dist < 0.5 * self.lines[i].height + self.viewport().setCursor(Qt.IBeamCursor if near else Qt.ArrowCursor) + else: + self.viewport().setCursor(Qt.OpenHandCursor if self.space else Qt.ArrowCursor) + super().mouseMoveEvent(ev) + + def mouseReleaseEvent(self, ev): + if self.pan_from is not None: + self.pan_from = None + self.viewport().setCursor(Qt.OpenHandCursor if self.space else Qt.ArrowCursor) + super().mouseReleaseEvent(ev) + + def keyPressEvent(self, ev): + if ev.key() == Qt.Key_Space and not ev.isAutoRepeat(): + self.space = True + self.viewport().setCursor(Qt.OpenHandCursor) + return + super().keyPressEvent(ev) + + def keyReleaseEvent(self, ev): + if ev.key() == Qt.Key_Space and not ev.isAutoRepeat(): + self.space = False + self.viewport().setCursor(Qt.ArrowCursor) + return + super().keyReleaseEvent(ev) + + def contextMenuEvent(self, ev): + if not self.lines: + return + m = QMenu(self) + m.addAction("Copy", self.copy).setEnabled(self.sel[0] != self.sel[1]) + m.addAction("Copy all text", lambda: QGuiApplication.clipboard().setText(self.text)) + m.addSeparator() + m.addAction("Select all", self.select_all) + m.exec(ev.globalPos()) + + def copy(self): + if self.selected_text(): + QGuiApplication.clipboard().setText(self.selected_text()) + + # ---- zoom ---- def fit(self): self.auto_fit = True @@ -274,7 +409,8 @@ class MainWindow(QMainWindow): self.setAcceptDrops(True) self.opts = opts self.ocr = Ocr() - self.image, self.path, self.lines, self.spans = None, None, [], [] + self.image, self.path, self.lines, self.text = None, None, [], "" + self.page_title = "Page" self.ocr_thread, self.scan_thread, self.scanning = None, None, False self.bridge = Bridge() self.bridge.frame.connect(self.on_frame) @@ -318,7 +454,7 @@ class MainWindow(QMainWindow): self.act_ocr = act("Read text", self.run_ocr, "Ctrl+R", "Run OCR on the page") self.act_auto = act("Auto", lambda: None, tip="Read the text after every scan", checkable=True) self.act_auto.setChecked(True) - self.act_copy = act("Copy text", self.copy_text, "Ctrl+Shift+C", "Copy all text to the clipboard") + self.act_copy = act("Copy all", self.copy_text, "Ctrl+Shift+C", "Copy all text to the clipboard") self.act_savetxt = act("Save text…", self.save_text, "Ctrl+Shift+S", "Save the text as a .txt file") tb.addSeparator() self.act_boxes = act("Boxes", lambda on: self.view.show_boxes(on), "Ctrl+B", "Show recognised lines on the page", True) @@ -331,22 +467,19 @@ class MainWindow(QMainWindow): self.camera = ImageLabel("camera") self.view = PageView() - self.view.lineClicked.connect(self.select_line) - self.text = QPlainTextEdit() - self.text.setPlaceholderText("Recognised text appears here; scan a page, or open or drop an image.") - self.text.setStyleSheet("font-size: 14px;") - self.text.setLineWrapMode(QPlainTextEdit.NoWrap) # one OCR line per text line - self.text.cursorPositionChanged.connect(self.on_cursor) - self.text.textChanged.connect(self.on_text_edited) - self._setting_text = False + self.view.selectionChanged.connect(self.on_selection) + for key, slot in ((QKeySequence.Copy, self.view.copy), (QKeySequence.SelectAll, self.view.select_all)): + a = QAction(self) + a.setShortcut(key) + a.triggered.connect(slot) + self.addAction(a) self.camera_pane, _ = pane("Camera", self.camera) page_pane, self.page_head = pane("Page", self.view) - text_pane, self.text_head = pane("Text", self.text) main = QSplitter(Qt.Horizontal) - for w in (self.camera_pane, page_pane, text_pane): + for w in (self.camera_pane, page_pane): main.addWidget(w) - main.setSizes([170, 750, 480]) + main.setSizes([220, 1180]) self.setCentralWidget(main) self.status_label = QLabel() self.statusBar().addWidget(self.status_label, 1) @@ -385,8 +518,7 @@ class MainWindow(QMainWindow): self.scanning = True self.view.set_lines([]) self.view.set_image(QImage()) - self.set_text("", []) - self.lines, self.image, self.path = [], None, None + self.lines, self.text, self.image, self.path = [], "", None, None for a in (self.act_save, self.act_ocr, self.act_copy, self.act_savetxt): a.setEnabled(False) self.page_head.setText("Page (scanning)") @@ -415,8 +547,12 @@ class MainWindow(QMainWindow): # ---- images ---- def load(self, path, auto_ocr=True): - img = imread(path) + try: + img = imread(path) + except OSError: + img = None if img is None: + log_error(f"cannot read {path}") QMessageBox.warning(self, "myScan", f"cannot read {path}") return self.set_page(img, path) @@ -424,12 +560,12 @@ class MainWindow(QMainWindow): self.run_ocr() def set_page(self, img, path=None): - self.image, self.path, self.lines = img, path, [] + self.image, self.path, self.lines, self.text = img, path, [], "" self.view.set_lines([]) self.view.set_image(to_qimage(img)) - self.set_text("", []) name = os.path.basename(path) if path else "page" - self.page_head.setText(f"Page — {name} ({img.shape[1]}×{img.shape[0]})") + self.page_title = f"Page — {name} ({img.shape[1]}×{img.shape[0]})" + self.page_head.setText(self.page_title) self.setWindowTitle(f"myScan — {name}") self.act_save.setEnabled(True) self.act_ocr.setEnabled(True) @@ -445,7 +581,12 @@ class MainWindow(QMainWindow): base = os.path.splitext(self.path)[0] if self.path else os.path.join(self.default_dir(), "page") path, _ = QFileDialog.getSaveFileName(self, "Save image", base + ".png", IMAGE_FILTER) if path: - imwrite(path, self.image) + try: + imwrite(path, self.image) + except OSError as e: + log_error(f"cannot save {path}: {e}") + QMessageBox.warning(self, "myScan", f"cannot save {path}: {e}") + return self.set_message(f"saved {path}") def default_dir(self): @@ -469,7 +610,7 @@ class MainWindow(QMainWindow): if self.image is None or (self.ocr_thread and self.ocr_thread.isRunning()): return self.act_ocr.setEnabled(False) - self.text_head.setText("Text — reading…") + self.page_head.setText(f"{self.page_title} · reading…") self.set_message("reading text…") self.ocr_thread = OcrThread(self.ocr, self.image) self.ocr_thread.done.connect(self.on_ocr) @@ -479,52 +620,32 @@ class MainWindow(QMainWindow): def on_ocr(self, lines, secs): self.act_ocr.setEnabled(self.image is not None) self.lines = lines - text, spans = layout(lines) - self.view.set_lines(lines) - self.set_text(text, spans) + self.text, spans = layout(lines) + self.view.set_lines(lines, self.text, spans) conf = f", mean confidence {np.mean([ln.score for ln in lines]):.0%}" if lines else "" - self.text_head.setText(f"Text — {len(lines)} lines{conf}") - self.set_message(f"read {len(lines)} lines in {secs:.1f} s") - self.act_copy.setEnabled(bool(text)) - self.act_savetxt.setEnabled(bool(text)) + self.page_head.setText(f"{self.page_title} · {len(lines)} lines{conf}") + self.set_message(f"read {len(lines)} lines in {secs:.1f} s" + (" · drag over the text to select it" + if lines else "")) + self.act_copy.setEnabled(bool(self.text)) + self.act_savetxt.setEnabled(bool(self.text)) def on_ocr_failed(self, msg): self.act_ocr.setEnabled(self.image is not None) - self.text_head.setText("Text") + self.page_head.setText(self.page_title) self.set_message(f"OCR failed: {msg}") - def set_text(self, text, spans): - self._setting_text = True - self.text.setPlainText(text) - self._setting_text = False - self.spans = spans - - def on_text_edited(self): - if not self._setting_text and self.spans: # positions no longer match the boxes - self.spans = [] - self.view.highlight(None) - self.act_copy.setEnabled(bool(self.text.toPlainText())) - self.act_savetxt.setEnabled(bool(self.text.toPlainText())) - - def on_cursor(self): - pos = self.text.textCursor().position() - hit = next((i for i, (a, b) in enumerate(self.spans) if a <= pos <= b), None) - if hit is not None or not self.spans: - self.view.highlight(hit) - - def select_line(self, i): - if i >= len(self.spans): + def on_selection(self, text): + if not text: return - a, b = self.spans[i] - c = self.text.textCursor() - c.setPosition(a) - c.setPosition(b, QTextCursor.KeepAnchor) - self.text.setTextCursor(c) - self.text.setFocus() - self.view.highlight(i) + cb = QGuiApplication.clipboard() + if cb.supportsSelection(): # X11/Wayland: middle-click pastes it elsewhere + cb.setText(text, QClipboard.Selection) + preview = " ⏎ ".join(t for t in text.split("\n") if t) + preview = preview if len(preview) <= 90 else preview[:87] + "…" + self.set_message(f"{len(text)} characters selected (Ctrl+C to copy): {preview}") def copy_text(self): - QGuiApplication.clipboard().setText(self.text.toPlainText()) + QGuiApplication.clipboard().setText(self.text) self.set_message("text copied to the clipboard") def save_text(self): @@ -532,7 +653,7 @@ class MainWindow(QMainWindow): path, _ = QFileDialog.getSaveFileName(self, "Save text", base + ".txt", "Text (*.txt)") if path: with open(path, "w", encoding="utf-8") as f: - f.write(self.text.toPlainText() + "\n") + f.write(self.text + "\n") self.set_message(f"saved {path}") def closeEvent(self, ev): @@ -557,6 +678,16 @@ def main(argv=None, opts=None): app.setApplicationName("myScan") w = MainWindow(a) w.show() + # Ctrl+C in the terminal closes the window (which puts the mouse back in mouse mode). + # Python only runs signal handlers between bytecodes, so wake it up now and then. + def interrupt(*_): + while (dialog := QApplication.activeModalWidget()) is not None: + dialog.reject() # a dialog's event loop would keep the app alive + dialog.hide() + w.close() + signal.signal(signal.SIGINT, interrupt) + tick = QTimer(timeout=lambda: None, interval=200) + tick.start() if a.image: - w.load(a.image) + QTimer.singleShot(0, lambda: w.load(a.image)) return app.exec() diff --git a/myscan/ocr.py b/myscan/ocr.py index 1ac1311..fd7bbf9 100644 --- a/myscan/ocr.py +++ b/myscan/ocr.py @@ -5,8 +5,12 @@ multilingual PP-OCRv6 models ship inside the wheel, so it works offline. They re Latin script (English, Swedish, German, French, ...) as well as Chinese and Japanese. Independent of the scanner hardware, like myscan.stitching. - lines = Ocr()(bgr) # [Line(box, text, score)], box = 4x2 corner points + lines = Ocr()(bgr) # [Line(box, text, score, edges)], box = 4x2 corner points text, spans = layout(lines) # plain text; spans[i] = (start, end) of lines[i] in it + +Each line knows where its characters are (from the recogniser's per-character boxes), +so text can be selected on the image: Line.locate maps a point to a character +position, Line.quad gives the outline of a range of characters. """ import threading from dataclasses import dataclass @@ -19,6 +23,11 @@ class Line: box: np.ndarray # 4x2 float corners (clockwise from top left), image pixels text: str score: float + edges: np.ndarray = None # len(text) + 1 character boundaries, 0..1 along the line + + def __post_init__(self): + if self.edges is None: # no character boxes: assume evenly spaced characters + self.edges = np.linspace(0.0, 1.0, len(self.text) + 1) @property def cy(self): @@ -34,6 +43,67 @@ class Line: grows with the slant of a skewed page).""" return float(np.hypot(*(self.box[3] - self.box[0]))) + def local(self, p): + """Image point -> (u, v): u runs 0..1 along the line, v 0..1 from its top to + its bottom edge.""" + b = self.box + m = np.array([b[1] - b[0], b[3] - b[0]]).T + return np.linalg.solve(m, np.asarray(p, float) - b[0]) + + def locate(self, p): + """(distance in pixels from the line's box, character position nearest p).""" + u, v = self.local(p) + b = self.box + du = max(-u, u - 1, 0) * np.hypot(*(b[1] - b[0])) + dv = max(-v, v - 1, 0) * self.height + return float(np.hypot(du, dv)), int(np.abs(self.edges - u).argmin()) + + def quad(self, a, b): + """4x2 outline of characters a..b (positions in text).""" + top, bot = self.box[[0, 1]], self.box[[3, 2]] + ta, tb = self.edges[a], self.edges[b] + lerp = lambda e, t: e[0] + t * (e[1] - e[0]) + return np.array([lerp(top, ta), lerp(top, tb), lerp(bot, tb), lerp(bot, ta)]) + + def word(self, k): + """(start, end) of the word around character position k.""" + a = b = min(k, len(self.text)) + while a > 0 and not self.text[a - 1].isspace(): + a -= 1 + while b < len(self.text) and not self.text[b].isspace(): + b += 1 + return a, b + + +def char_edges(box, text, chars): + """Character boundaries along a line (0..1) from the recogniser's boxes for its + non-space characters (4x2 corners each); a space spans the gap between its + neighbours. None if the boxes don't match the text.""" + if not chars or len(chars) != len(text) - text.count(" "): + return None + m0, m1 = (box[0] + box[3]) / 2, (box[1] + box[2]) / 2 + d = m1 - m0 + t = lambda p: float(np.dot(p - m0, d) / np.dot(d, d)) + ext = iter([(t((c[0] + c[3]) / 2), t((c[1] + c[2]) / 2)) for c in (np.asarray(c, float) for c in chars)]) + lo, hi = [None] * len(text), [None] * len(text) + for i, ch in enumerate(text): + if ch != " ": + lo[i], hi[i] = next(ext) + end = 0.0 + for i in range(len(text)): # spaces start where the previous character ends + if lo[i] is None: + lo[i] = end + else: + end = hi[i] + start = 1.0 + for i in reversed(range(len(text))): # ... and end where the next one starts + if hi[i] is None: + hi[i] = start + else: + start = lo[i] + edges = [0.0] + [(hi[i - 1] + lo[i]) / 2 for i in range(1, len(text))] + [1.0] + return np.clip(np.maximum.accumulate(edges), 0.0, 1.0) + class Ocr: """Lazily loads the engine (about 0.3 s) on first use; calls are serialised.""" @@ -50,10 +120,15 @@ class Ocr: def __call__(self, bgr): with self._lock: - r = self.engine()(bgr) + r = self.engine()(bgr, return_word_box=True, return_single_char_box=True) if r.boxes is None: return [] - return [Line(np.asarray(b, float), t, float(s)) for b, t, s in zip(r.boxes, r.txts, r.scores)] + words = r.word_results or [None] * len(r.txts) + out = [] + for b, t, s, w in zip(r.boxes, r.txts, r.scores, words): + b = np.asarray(b, float) + out.append(Line(b, t, float(s), char_edges(b, t, [c[2] for c in w]) if w else None)) + return out def rows(lines):