95cab33f19
Merge in RR/cv-analysis from evaluate_layout_detection to master
Squashed commit of the following:
commit 8ec2f69fc61d1e15bd502b0a2c1f720cbec2b34e
Author: llocarnini <lillian.locarnini@iqser.com>
Date: Tue Aug 23 15:07:21 2022 +0200
repaired is_not_included() logic (did drop the outer rectangle, not the included)
commit 97be081d1e60989313924ceac0bfb3062229411e
Merge: 2c28fa2 2b5c4f1
Author: llocarnini <lillian.locarnini@iqser.com>
Date: Tue Aug 23 14:28:14 2022 +0200
Merge branch 'master' of ssh://git.iqser.com:2222/rr/cv-analysis into evaluate_layout_detection
commit 2c28fa280b7eff922c715245fffe69702c7e6742
Author: llocarnini <lillian.locarnini@iqser.com>
Date: Tue Aug 23 13:50:17 2022 +0200
del print statements
commit c60121fc4faebc5de556ec0ab7a3af4f815f7ce1
Author: llocarnini <lillian.locarnini@iqser.com>
Date: Mon Aug 22 10:51:52 2022 +0200
few changes to connect_rects.py
commit a99719905d58cbe856fa020177abd7e317c1d072
Author: llocarnini <lillian.locarnini@iqser.com>
Date: Thu Aug 18 08:37:12 2022 +0200
layout parsing improved with connect_rects.py
commit d693688a0f0d63395cfd36645de7b3417f64de30
Author: llocarnini <lillian.locarnini@iqser.com>
Date: Tue Aug 2 09:31:19 2022 +0200
removed vizlogger instances
64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import argparse
|
|
import timeit
|
|
from time import process_time
|
|
from itertools import starmap
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from funcy import lmap
|
|
|
|
from cv_analysis.figure_detection.figure_detection import detect_figures
|
|
from cv_analysis.layout_parsing import parse_layout
|
|
from cv_analysis.table_parsing import parse_tables
|
|
from cv_analysis.utils.draw import draw_rectangles
|
|
from pdf2img.conversion import convert_pages_to_images
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("pdf_path")
|
|
parser.add_argument("--output_folder", default="/tmp")
|
|
parser.add_argument("--type", "-t", choices=["table", "layout", "figure"], required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def analyse_and_annotate(images, analysis_fn):
|
|
arrays = lmap(lambda image: image.asarray(), images)
|
|
result = map(analysis_fn, arrays)
|
|
annotated_images = starmap(draw_rectangles, zip(arrays, result))
|
|
return annotated_images
|
|
|
|
|
|
def save_as_pdf(images, output_folder, file_name, operation):
|
|
Path(output_folder).mkdir(parents=True, exist_ok=True)
|
|
images = lmap(Image.fromarray, images)
|
|
images[0].save(f"{output_folder}/{file_name}_annotated_{operation}.pdf", save_all=True, append_images=images)
|
|
|
|
|
|
def get_analysis_fn(analysis_type):
|
|
if analysis_type == "table":
|
|
return parse_tables
|
|
elif analysis_type == "layout":
|
|
return parse_layout
|
|
elif analysis_type == "figure":
|
|
return detect_figures
|
|
else:
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
t0 = timeit.default_timer()
|
|
with open(args.pdf_path, "rb") as f:
|
|
pdf_bytes = f.read()
|
|
images = convert_pages_to_images(pdf_bytes)
|
|
t1 = timeit.default_timer()
|
|
annotated_pages = analyse_and_annotate(images=images, analysis_fn=get_analysis_fn(args.type))
|
|
t2 = timeit.default_timer()
|
|
save_as_pdf(annotated_pages, args.output_folder, Path(args.pdf_path).stem, args.type)
|
|
t3 = timeit.default_timer()
|
|
print("[s] opening file and convert pdf pages to images: ", t1-t0)
|
|
print("[s] analyse and annotate images: ", t2-t1)
|
|
print("[s] save images as pdf: ", t3-t2)
|
|
|