Marmelator/image-classification-service
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 29 Wiki Activity

implement ad hoc channel count detection for new image extraction

Browse Source
This commit is contained in:
Julius Unverfehrt 2023-02-02 13:57:56 +01:00
parent 73f7491c8f
commit c0b41e77b8
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
image_prediction/image_extractor/extractors/parsable.py
View File

@ -136,7 +136,11 @@ 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)
array = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(pixmap.h, pixmap.w)
# 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)
return Image.fromarray(array)