diff --git a/dlibrary/dlibrary.py b/dlibrary/dlibrary.py index acc3291..fcaf440 100755 --- a/dlibrary/dlibrary.py +++ b/dlibrary/dlibrary.py @@ -77,6 +77,8 @@ 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") @@ -232,7 +234,7 @@ def collate(args): work_staging_dir = collation_staging_area / work_id - collator = Collator(work_staging_dir, [], args.locale) + collator = Collator(work_staging_dir, [], args) 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}') @@ -252,10 +254,10 @@ def collate(args): con.close() class Collator: - def __init__(self, dest, exclude, locale): + def __init__(self, dest, exclude, args): self.dest = dest self.exclude = exclude - self.locale = locale + self.args = args self.index = 0 def collate_from_paths(self, srcs): @@ -309,14 +311,13 @@ class Collator: def link_pdf(self, src): with fitz.open(src) as pdf: - xrefs = image_xrefs(pdf) - if xrefs is None: + images = pdf_images(pdf, self.args.force_convert_pdf) + if images 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, xref) in enumerate(xrefs, start=self.index): - image = pdf.extract_image(xref) + for (idx, image) in enumerate(images, start=self.index): file_path = self.dest / f'{idx:04d}.{image["ext"]}' with open(file_path, 'wb') as f: f.write(image["image"]) @@ -422,34 +423,53 @@ class Collator: return False def try_collate_select_language(self, srcs): - if self.locale not in LANGUAGE_REGEXES: + if self.args.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.locale].search(nname(src))] + srcs_matching_language = [src for src in srcs if LANGUAGE_REGEXES[self.args.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 image_xrefs(pdf): +def pdf_images(pdf, force=False): 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] + return (pdf.extract_image(images[0][0]) for images in images_by_page) print("Checking PDF images the quick way failed, trying the slow way") - 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']) + 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: + 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('') - print('\nSuccess') - return xrefs + 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) def nfc(s): return unicodedata.normalize('NFC', s) @@ -706,7 +726,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.locale) + collator = Collator(work_staging_dir, exclusions, args) 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: @@ -1002,7 +1022,7 @@ parser_collate.add_argument( nargs='*', help='paths within extraction folders as collation hints' ) -parser_collate.set_defaults(func=collate) +parser_collate.set_defaults(func=collate, force_convert_pdf=False) parser_manual_collate = subparsers.add_parser( 'manual-collate', @@ -1042,6 +1062,11 @@ 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='+',