alpha channel info WIP

image_prediction/image_extractor/extractors/parsable.py

@@ -42,7 +42,7 @@ class ParsablePDFImageExtractor(ImageExtractor):
 
     def __process_images_on_page(self, page: fitz.fitz.Page):
         images = get_images_on_page(self.doc, page)
-        metadata = get_metadata_for_images_on_page(page)
+        metadata = get_metadata_for_images_on_page(self.doc, page)
         get_image_infos.cache_clear()
         load_image_handle_from_xref.cache_clear()
 
@@ -69,18 +69,21 @@ def get_images_on_page(doc, page: fitz.Page):
     return images
 
 
-def get_metadata_for_images_on_page(page: fitz.Page):
+def get_metadata_for_images_on_page(doc, page: fitz.Page):
     image_infos = get_image_infos(page)
+
     metadata = map(get_image_metadata, image_infos)
     metadata = validate_coords_and_passthrough(metadata)
+
     metadata = filterfalse(tiny, metadata)
     metadata = validate_size_and_passthrough(metadata)
+
     metadata = map(partial(merge, get_page_metadata(page)), metadata)
 
-    # xrefs = map(itemgetter("xref"), image_infos)
-    # alpha = map(has_alpha_channel, xrefs)
-    # alpha = zipdict(repeat(Info.ALPHA), alpha)
-    # metadata = starmap(merge, zip(alpha, metadata))
+    xrefs = map(itemgetter("xref"), image_infos)
+    alpha = map(partial(has_alpha_channel, doc), xrefs)
+    alpha = ({Info.ALPHA: a} for a in alpha)
+    metadata = starmap(merge, zip(alpha, metadata))
 
     yield from metadata
 

image_prediction/info.py

@@ -11,4 +11,4 @@ class Info(Enum):
     X2 = "x2"
     Y1 = "y1"
     Y2 = "y2"
-    # ALPHA = "alpha"
+    ALPHA = "alpha"

test/conftest.py

@@ -186,8 +186,18 @@ def batch_size(request):
     return request.param
 
 
+@pytest.fixture(params=[False])
+def input_size(request, __input_size):
+    alpha = request.param
+    print(alpha)
+    if alpha:
+        w, h, d = __input_size
+        __input_size = w, h, d + 1
+    return __input_size
+
+
 @pytest.fixture(params=[{"width": 10, "height": 15, "depth": 3}, {"width": 150, "height": 100, "depth": 3}])
-def input_size(request):
+def __input_size(request):
     return itemgetter("width", "height", "depth")(request.param)
 
 
@@ -291,6 +301,7 @@ def metadata(images, info_label_map):
             info_label_map.X2: x2,
             info_label_map.Y1: y1,
             info_label_map.Y2: y2,
+            info_label_map.ALPHA: image.mode == "RGBA"
         }
         return metadata