diff --git a/dlibrary/dlibrary.py b/dlibrary/dlibrary.py index 864d6b6..acc3291 100755 --- a/dlibrary/dlibrary.py +++ b/dlibrary/dlibrary.py @@ -77,8 +77,6 @@ IMAGE_FILE_EXTENSIONS = ['.png', '.jpg', '.jpeg', '.gif', '.tiff', '.bmp'] IGNOREABLE_FILES = ['Thumbs.db', '__MACOSX', '.DS_Store'] IGNOREABLE_EXTENSIONS = ['.txt', '.html', '.htm', '.psd', '.mp4'] -PDF_FALLBACK_DPI = 300 - def open_zipfile_with_encoding(path): try: return zipfile.ZipFile(path, metadata_encoding="utf-8") @@ -208,15 +206,12 @@ def fetch(args): asyncio.run(fetch_async(args)) -def self_and_parents(path): - return [path] + list(path.parents) - def collate(args): con = sqlite3.connect(args.destdir / 'meta.db') cur = con.cursor() extraction_dir = args.destdir / 'extract' - hint_map = {self_and_parents(Path(relpath(hint, extraction_dir)))[-2].name: hint for hint in args.hints} + hint_map = {Path(relpath(hint, extraction_dir)).parents[-2].name: hint for hint in args.hints} collation_staging_area = args.destdir / 'site' / 'images-staging' collation_staging_area.mkdir(parents=True) @@ -237,7 +232,7 @@ def collate(args): work_staging_dir = collation_staging_area / work_id - collator = Collator(work_staging_dir, [], args) + collator = Collator(work_staging_dir, [], args.locale) collation_result = collator.collate_from_paths([hint_map.get(work_id, work_path)]) if collation_result and collator.index > 0: print(f'Collated {collator.index} pages for {work_id}') @@ -257,10 +252,10 @@ def collate(args): con.close() class Collator: - def __init__(self, dest, exclude, args): + def __init__(self, dest, exclude, locale): self.dest = dest self.exclude = exclude - self.args = args + self.locale = locale self.index = 0 def collate_from_paths(self, srcs): @@ -314,13 +309,14 @@ class Collator: def link_pdf(self, src): with fitz.open(src) as pdf: - images = pdf_images(pdf, self.args.force_convert_pdf) - if images is None: + xrefs = image_xrefs(pdf) + if xrefs is None: print(f'Support for weirder PDFs not yet implemented, skipping {src}') return None self.dest.mkdir(parents=True, exist_ok=True) - for (idx, image) in enumerate(images, start=self.index): + for (idx, xref) in enumerate(xrefs, start=self.index): + image = pdf.extract_image(xref) file_path = self.dest / f'{idx:04d}.{image["ext"]}' with open(file_path, 'wb') as f: f.write(image["image"]) @@ -426,57 +422,34 @@ class Collator: return False def try_collate_select_language(self, srcs): - if self.args.locale not in LANGUAGE_REGEXES: + if self.locale not in LANGUAGE_REGEXES: return False if not all(any(lang.search(nname(src)) for lang in LANGUAGE_REGEXES.values()) for src in srcs): return False - srcs_matching_language = [src for src in srcs if LANGUAGE_REGEXES[self.args.locale].search(nname(src))] + srcs_matching_language = [src for src in srcs if LANGUAGE_REGEXES[self.locale].search(nname(src))] if len(srcs_matching_language) == len(srcs) or len(srcs_matching_language) == 0: return False return self.collate_from_paths(srcs_matching_language) -def is_single_image(page): - blocks = page.get_text('blocks') - return len(blocks) == 1 and blocks[0][6] == 1 - -def pdf_images(pdf, force=False): - images_by_page = [(page.get_images(), is_single_image(page)) for page in pdf] - if all(len(images) == 1 and single for (images, single) in images_by_page): - return (pdf.extract_image(images[0][0]) for (images, _) in images_by_page) +def image_xrefs(pdf): + images_by_page = [page.get_images() for page in pdf] + if all(len(images) == 1 for images in images_by_page): + return [images[0][0] for images in images_by_page] print("Checking PDF images the quick way failed, trying the slow way") - def xref_or_image_generator(): - xref_mode = not force - for (idx, page) in enumerate(pdf): - page_images = page.get_image_info(xrefs=True) - if len(page_images) == 1 and page_images[0]['xref'] != 0 and is_single_image(page): - xref = page_images[0]['xref'] - if xref_mode: - yield xref - else: - yield pdf.extract_image(xref) - else: - if xref_mode: - raise ValueError - else: - print(f'\nGenerating pixmap for page {idx+1}') - pix = page.get_pixmap(dpi=PDF_FALLBACK_DPI) - yield { 'ext': 'png', 'image': pix.tobytes('png') } - print(f'\x1b[2K\r{idx+1}/{pdf.page_count} pages processed...', end='') - print('') + xrefs = [] + for (idx, page) in enumerate(pdf): + print(f'\x1b[2K\r{idx}/{pdf.page_count} pages processed...', end='') + images = page.get_image_info(xrefs=True) + if len(images) != 1 or images[0]['xref'] == 0: + print('\nFailed') + return None + xrefs.append(images[0]['xref']) - if force: - return xref_or_image_generator() - - try: - xrefs = list(xref_or_image_generator()) - except ValueError: - print('\nFailed') - return None - print('Success') - return (pdf.extract_image(xref) for xref in xrefs) + print('\nSuccess') + return xrefs def nfc(s): return unicodedata.normalize('NFC', s) @@ -704,7 +677,7 @@ def manual_collate(args): extraction_dir = args.destdir / 'extract' sample_path = next(path for group in (raw_groups + [raw_exclusions]) for path in group) - work_id = self_and_parents(Path(relpath(sample_path, extraction_dir)))[-2].name + work_id = Path(relpath(sample_path, extraction_dir)).parents[-2].name exclusions = [normalize_to(item, args.destdir) for item in raw_exclusions] @@ -733,7 +706,7 @@ def manual_collate(args): work_staging_dir = collation_staging_area / work_id work_staging_dir.mkdir(parents=True) - collator = Collator(work_staging_dir, exclusions, args) + collator = Collator(work_staging_dir, exclusions, args.locale) for group in groups: collation_result = collator.collate_from_paths([item for item in group if item not in exclusions]) if collation_result is None: @@ -1029,7 +1002,7 @@ parser_collate.add_argument( nargs='*', help='paths within extraction folders as collation hints' ) -parser_collate.set_defaults(func=collate, force_convert_pdf=False) +parser_collate.set_defaults(func=collate) parser_manual_collate = subparsers.add_parser( 'manual-collate', @@ -1069,11 +1042,6 @@ parser_manual_collate = subparsers.add_parser( only handle one work at a time. """), ) -parser_manual_collate.add_argument( - '--force-convert-pdf', - action='store_true', - help="convert a PDF page to a 300dpi image if there isn't a single image we can extract directly", -) parser_manual_collate.add_argument( 'expression', nargs='+',