1#!/usr/bin/env python3
2"""
3Get OCR'd text for all the images in a directory using Live Text.
4
5See https://alexwlchan.net/2022/live-text-script/
6
7== Usage ==
8
9Pass the name of the directory you want to scan as a single argument:
10
11    $ python3 get_all_live_text ~/screenshots
12
13It will create a JSON file `live_text.json` in the top of the directory with
14the text for every image it finds.  So the command above would create a file
15at `~/screenshots/live_text.json`.
16
17There will be one line for each image, for example:
18
19    {"path": "railway-sign.jpg", "text": ["Passengers must", "not pass this point", "or cross the line"]}
20    {"path": "dancers.jpg", "text": []}
21
22The script runs incrementally, so if you add more images later, you can
23re-run to just get the text for all the new images.
24
25"""
26
27import json
28import os
29import subprocess
30import sys
31
32try:
33    from tqdm import tqdm
34except ImportError:
35    def tqdm(s):
36        return s
37
38
39def get_file_paths_under(root):
40    """Generates the paths to every file under ``root``."""
41    if not os.path.isdir(root):
42        raise ValueError(f"Cannot find files under non-existent directory: {root!r}")
43
44    for dirpath, _, filenames in os.walk(root):
45        for f in filenames:
46            if os.path.isfile(os.path.join(dirpath, f)):
47                yield os.path.join(dirpath, f)
48
49
50def get_static_image_paths_under(root):
51    for path in get_file_paths_under(root):
52        if os.path.basename(path).startswith("."):
53            continue
54        if path.lower().endswith((".gif", ".pdf", ".webarchive")):
55            continue
56        yield path
57
58
59def get_text(path):
60    try:
61        return json.loads(subprocess.check_output(["get_live_text", path]))
62    except subprocess.CalledProcessError:
63        return None
64
65
66if __name__ == "__main__":
67    try:
68        root = sys.argv[1]
69    except IndexError:
70        sys.exit(f"Usage: {__file__} <PATH>")
71
72    out_path = os.path.join(root, "live_text.json")
73
74    try:
75        seen_paths = {json.loads(line)["path"] for line in open(out_path)}
76    except FileNotFoundError:
77        seen_paths = set()
78
79    for path in tqdm(list(get_static_image_paths_under(root))):
80        if path in seen_paths:
81            continue
82
83        with open(out_path, "a") as outfile:
84            outfile.write(json.dumps({"path": path, "text": get_text(path)}) + "\n")