From c4416636c0a72b879e9b8d511f791a5e00b301fa Mon Sep 17 00:00:00 2001 From: Julius Unverfehrt Date: Thu, 2 Feb 2023 14:10:32 +0100 Subject: [PATCH] beautify --- image_prediction/image_extractor/extractors/parsable.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/image_prediction/image_extractor/extractors/parsable.py b/image_prediction/image_extractor/extractors/parsable.py index 941f431..a847e17 100644 --- a/image_prediction/image_extractor/extractors/parsable.py +++ b/image_prediction/image_extractor/extractors/parsable.py @@ -136,11 +136,8 @@ def get_image_infos(page: fitz.Page) -> List[dict]: @lru_cache(maxsize=None) def xref_to_image(doc, xref) -> Image: pixmap = fitz.Pixmap(doc, xref) - # FIXME: implement proper logic to determine how many channels the image has - try: - array = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(pixmap.h, pixmap.w) - except ValueError: - array = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(pixmap.h, pixmap.w, pixmap.n) + array = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(pixmap.h, pixmap.w, pixmap.n) + array = array[:, :, 0] if array.shape[2] == 1 else array return Image.fromarray(array)