ocrseals/ocr_seals.py at main · idsinge/ocrseals · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import argparse
import os, sys, subprocess
import shlex, glob
from pathlib import Path

import classify_and_transribe

def main(args):

    input_dir = os.path.join(args.data, 'input')
    input_files = os.listdir(input_dir)

    for infile in input_files:

        instem = Path(infile).stem

        if not Path(infile).suffix in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']: continue

        try:
            os.remove( os.path.join(args.data, "detect", "labels", f"{instem}.txt") )
            map(os.remove, glob.glob(f"{args.data}/detect/crops/char/{instem}*"))
        except OSError:
            pass

    detect_cmd = f"python3 detection/yolov5/detect.py --weights {args.detection_model} --source {input_dir} --imgsz 1000  --save-crop  --save-txt --project {args.data} --name detect --exist-ok --device cpu"

    print("\n ===== Applying detection =====\n\n")

    subprocess.run(shlex.split(detect_cmd))


    print("\n ===== Applying Hough line detection =====\n\n")

    for infile in input_files:

        instem = Path(infile).stem

        if not Path(infile).suffix in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']: continue

        hough_cmd = f"python3 hough/Hough.py --PATH_data {input_dir}/ --ann_name '../detect/labels/{instem}.txt' --img_name '{infile}'  --save_dir {args.data}/hough/"

        print(f"-- {hough_cmd} --\n")

        subprocess.run(shlex.split(hough_cmd))


    print("\n ===== Applying classification and extracting transcription =====\n\n")

    cmodel = classify_and_transribe.load_model(args.classification_model)

    os.makedirs(os.path.join(args.data, "trans"), exist_ok=True)

    for infile in input_files:

        instem = Path(infile).stem

        if not Path(infile).suffix in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']: continue

        boxes_path = os.path.join(args.data, "detect", "labels", f"{instem}.txt")
        crops_path = os.path.join(args.data, "detect", "crops", "char")
        keys_path = os.path.join(args.data, "hough", f"{infile}.keys.txt")

        outfile = os.path.join(args.data, "trans", f"{instem}.pred.txt")

        print(f"-- Processing {infile} and saving trnascription to {outfile} --\n")

        classified_boxes = classify_and_transribe.classify(cmodel, boxes_path, crops_path)
        classify_and_transribe.transcript(classified_boxes, keys_path, outfile)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--detection-model', required=True, help="trained model to use for detection")
    parser.add_argument('--classification-model', required=True, help="trained model to use for classification")

    parser.add_argument('--data', required=True, help="data directory for inputs and output - inputs will be read from data/input")

    args = parser.parse_args()


    if os.path.isfile(args.detection_model) and os.path.isfile(args.classification_model) and os.path.isdir(args.data) and os.path.isdir(os.path.join(args.data, "input")):

        main(args)

    else:
        parser.print_usage()
        exit(1)