deepdoctection 0.26__py3-none-any.whl → 0.27__py3-none-any.whl

deepdoctection/utils/viz.py

@@ -24,17 +24,28 @@ and
 
 <https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/colormap.py>
 """
+
+import base64
+import os
 import sys
-from typing import Callable, List, Optional, Tuple, no_type_check
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Sequence, Tuple, no_type_check
 
-import cv2
 import numpy as np
 import numpy.typing as npt
-from numpy import float32
+from numpy import float32, uint8
 
 from .detection_types import ImageType
+from .file_utils import get_opencv_requirement, get_pillow_requirement, opencv_available, pillow_available
+
+if opencv_available():
+    import cv2
+
+if pillow_available():
+    from PIL import Image, ImageDraw
 
-__all__ = ["draw_text", "draw_boxes", "interactive_imshow"]
+
+__all__ = ["draw_boxes", "interactive_imshow", "viz_handler"]
 
 _COLORS = (
     np.array(
@@ -183,39 +194,6 @@ def random_color(rgb: bool = True, maximum: int = 255) -> Tuple[int, int, int]:
     return tuple(int(x) for x in ret)  # type: ignore
 
 
-def draw_text(
-    np_image: ImageType, pos: Tuple[int, int], text: str, color: Tuple[int, int, int], font_scale: float = 1.0
-) -> ImageType:
-    """
-    Draw text on an image.
-
-    :param np_image: image as np.ndarray
-    :param pos: x, y; the position of the text
-    :param text: text string to draw
-    :param color: a 3-tuple BGR color in [0, 255]
-    :param font_scale: float
-    :return: numpy array
-    """
-
-    np_image = np_image.astype(np.uint8)
-    x_0, y_0 = int(pos[0]), int(pos[1])
-    # Compute text size.
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    ((text_w, text_h), _) = cv2.getTextSize(text, font, font_scale, 1)
-    # Place text background.
-    if x_0 + text_w > np_image.shape[1]:
-        x_0 = np_image.shape[1] - text_w
-    if y_0 - int(1.15 * text_h) < 0:
-        y_0 = int(1.15 * text_h)
-    back_top_left = x_0, y_0 - int(1.3 * text_h)
-    back_bottom_right = x_0 + text_w, y_0
-    cv2.rectangle(np_image, back_top_left, back_bottom_right, color, -1)
-    # Show text.
-    text_bottomleft = x_0, y_0 - int(0.25 * text_h)
-    cv2.putText(np_image, text, text_bottomleft, font, font_scale, (0, 0, 0), thickness=1, lineType=cv2.LINE_AA)
-    return np_image
-
-
 def draw_boxes(
     np_image: ImageType,
     boxes: npt.NDArray[float32],
@@ -245,12 +223,13 @@ def draw_boxes(
         category_names = set(category_names_list)  # type: ignore
         category_to_color = {category: random_color() for category in category_names}
 
-    boxes = np.asarray(boxes, dtype="int32")
+    boxes = np.array(boxes, dtype="int32")
     if category_names_list is not None:
         assert len(category_names_list) == len(boxes), f"{len(category_names_list)} != {len(boxes)}"
     areas = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
     sorted_inds = np.argsort(-areas)  # draw large ones first
     assert areas.min() > 0, areas.min()
+
     # allow equal, because we are not very strict about rounding error here
     assert (
         boxes[:, 0].min() >= 0
@@ -270,71 +249,411 @@ def draw_boxes(
             if choose_color is None:
                 choose_color = random_color()
             if category_names_list[i] is not None:
-                np_image = draw_text(
+                np_image = viz_handler.draw_text(
                     np_image, (box[0], box[1]), category_names_list[i], color=choose_color, font_scale=font_scale
                 )
-            cv2.rectangle(
-                np_image, (box[0], box[1]), (box[2], box[3]), color=choose_color, thickness=rectangle_thickness
+            np_image = viz_handler.draw_rectangle(
+                np_image, (box[0], box[1], box[2], box[3]), choose_color, rectangle_thickness
             )
 
     # draw a (very ugly) color palette
     y_0 = np_image.shape[0]
     for category, col in category_to_color.items():
         if category is not None:
-            np_image = draw_text(np_image, (np_image.shape[1], y_0), category, color=col, font_scale=2)
-            font = cv2.FONT_HERSHEY_SIMPLEX
-            ((_, text_h), _) = cv2.getTextSize(category, font, font_scale, 2)
+            viz_handler.draw_text(np_image, (np_image.shape[1], y_0), category, color=col, font_scale=font_scale)
+            _, text_h = viz_handler.get_text_size(category, 2)
             y_0 = y_0 - int(10 * text_h)
 
     return np_image
 
 
 @no_type_check
-def interactive_imshow(
-    img: ImageType,
-    lclick_cb: Optional[Callable[[npt.NDArray[float32], int, int], None]] = None,
-    rclick_cb: Optional[Callable[[npt.NDArray[float32], int, int], None]] = None,
-    **kwargs: str,
-) -> None:
+def interactive_imshow(img: ImageType) -> None:
     """
     Display an image in a pop-up window
 
     :param img: An image (expect BGR) to show.
-    :param lclick_cb: a callback ``func(img, x, y)`` for left/right click event.
-    :param rclick_cb: a callback ``func(img, x, y)`` for left/right click event.
-    :param kwargs: can be {key_cb_a: callback_img, key_cb_b: callback_img}, to specify a callback ``func(img)`` for
-                   keypress. Some existing keypress event handler:
-
-                          * q: destroy the current window
-                          * x: execute ``sys.exit()``
-                          * s: save image to "out.png"
     """
-    name = "q, x: quit / s: save"
-    cv2.imshow(name, img)
-
-    def mouse_cb(event, x, y, *args):
-        if event == cv2.EVENT_LBUTTONUP and lclick_cb is not None:
-            lclick_cb(img.astype(dtype=float32), x, y)
-        elif event == cv2.EVENT_RBUTTONUP and rclick_cb is not None:
-            rclick_cb(img.astype(dtype=float32), x, y)
-
-    cv2.setMouseCallback(name, mouse_cb)
-    key = cv2.waitKey(-1)
-    while key >= 128:
+    viz_handler.interactive_imshow(img)
+
+
+class VizPackageHandler:
+    """
+    A handler for the image processing libraries PIL or OpenCV. Explicit use of the libraries is not intended.
+    If the environ.ment variable USE_OPENCV=True is set, only the CV2 functions will be used via the handler.
+    The default library is PIL. Compared to OpenCV, PIL is somewhat slower (this applies to reading and writing
+    image files), which can lead to a bottleneck during training, especially if the loading is not parallelized
+    """
+
+    PACKAGE_FUNCS = {
+        "cv2": {
+            "read_image": "_cv2_read_image",
+            "write_image": "_cv2_write_image",
+            "convert_np_to_b64": "_cv2_convert_np_to_b64",
+            "convert_b64_to_np": "_cv2_convert_b64_to_np",
+            "resize": "_cv2_resize",
+            "get_text_size": "_cv2_get_text_size",
+            "draw_rectangle": "_cv2_draw_rectangle",
+            "draw_text": "_cv2_draw_text",
+            "interactive_imshow": "_cv2_interactive_imshow",
+            "encode": "_cv2_encode",
+        },
+        "pillow": {
+            "read_image": "_pillow_read_image",
+            "write_image": "_pillow_write_image",
+            "convert_np_to_b64": "_pillow_convert_np_to_b64",
+            "convert_b64_to_np": "_pillow_convert_b64_to_np",
+            "resize": "_pillow_resize",
+            "get_text_size": "_pillow_get_text_size",
+            "draw_rectangle": "_pillow_draw_rectangle",
+            "draw_text": "_pillow_draw_text",
+            "interactive_imshow": "_pillow_interactive_imshow",
+            "encode": "_pillow_encode",
+        },
+    }
+
+    def __init__(self) -> None:
+        """Selecting the image processing library and fonts"""
+        package = self._select_package()
+        self.pkg_func_dict: Dict[str, str] = {}
+        self.font = None
+        self._set_vars(package)
+
+    @staticmethod
+    def _select_package() -> str:
+        """
+        USE_OPENCV has priority and will enforce to use OpenCV
+        Otherwise it will use Pillow as default package, if it is installed.
+        If Pillow is not installed it will try to load OpenCV again
+        :return: either 'pillow' or 'cv2'
+        """
+        if os.environ.get("USE_OPENCV") is not None:
+            requirements = get_opencv_requirement()
+            if not requirements[1]:
+                raise ImportError(requirements[2])
+            return "cv2"
+        requirements = get_pillow_requirement()
+        if os.environ["USE_PILLOW"]:
+            if not requirements[1]:
+                raise ImportError(requirements[2])
+            return "pillow"
+        requirements = get_opencv_requirement()
+        if not requirements[1]:
+            raise ImportError(requirements[2])
+        return "cv2"
+
+    def _set_vars(self, package: str) -> None:
+        self.pkg_func_dict = self.PACKAGE_FUNCS[package]
+        if package == "pillow":
+            image = Image.fromarray(np.uint8(np.ones((1, 1, 3))))
+            self.font = ImageDraw.ImageDraw(image).getfont()
+        else:
+            self.font = cv2.FONT_HERSHEY_SIMPLEX  # type: ignore
+
+    def refresh(self) -> None:
+        """
+        Refresh the viz_handler setting. Useful if you change the env variable on run time and want to take account of
+        the changes.
+
+        **Example**
+
+           os.env["USE_OPENCV"]="True"
+           viz_handler.refresh()             # this will reset the original config and now use OpenCV
+
+        :return:
+        """
+        package = self._select_package()
+        self._set_vars(package)
+
+    def read_image(self, path: str) -> ImageType:
+        """Reading an image from file and returning a np.array
+
+        :param path: Use /path/to/dir/file_name.[suffix]
+        """
+        return getattr(self, self.pkg_func_dict["read_image"])(path)
+
+    @staticmethod
+    def _cv2_read_image(path: str) -> ImageType:
+        return cv2.imread(path, cv2.IMREAD_COLOR)
+
+    @staticmethod
+    def _pillow_read_image(path: str) -> ImageType:
+        with Image.open(path) as image:
+            np_image = np.array(image)[:, :, ::-1]
+        return np_image
+
+    def write_image(self, path: str, image: ImageType) -> None:
+        """Writing an image as np.array to a file.
+
+        :param path: Use /path/to/dir/file_name.[suffix]
+        :param image: pixel values as np.array
+        """
+        return getattr(self, self.pkg_func_dict["write_image"])(path, image)
+
+    @staticmethod
+    def _cv2_write_image(path: str, image: ImageType) -> None:
+        cv2.imwrite(path, image)
+
+    @staticmethod
+    def _pillow_write_image(path: str, image: ImageType) -> None:
+        pil_image = Image.fromarray(np.uint8(image[:, :, ::-1]))
+        pil_image.save(path)
+
+    def encode(self, np_image: ImageType) -> bytes:
+        """Converting an image as np.array into a b64 representation
+
+        :param np_image: Image as np.array
+        """
+        return getattr(self, self.pkg_func_dict["encode"])(np_image)
+
+    @staticmethod
+    def _cv2_encode(np_image: ImageType) -> bytes:
+        np_encode = cv2.imencode(".png", np_image)
+        b_image = np_encode[1].tobytes()
+        return b_image
+
+    @staticmethod
+    def _pillow_encode(np_image: ImageType) -> bytes:
+        buffered = BytesIO()
+        pil_image = Image.fromarray(np.uint8(np_image[:, :, ::-1]))
+        pil_image.save(buffered, format="PNG")
+        return buffered.getvalue()
+
+    def convert_np_to_b64(self, image: ImageType) -> str:
+        """Converting an image given as np.array into a b64 encoded string
+
+        :param image: Image as np.array
+        """
+        return getattr(self, self.pkg_func_dict["convert_np_to_b64"])(image)
+
+    @staticmethod
+    def _cv2_convert_np_to_b64(image: ImageType) -> str:
+        np_encode = cv2.imencode(".png", image)
+        return base64.b64encode(np_encode[1]).decode("utf-8")  # type: ignore
+
+    @staticmethod
+    def _pillow_convert_np_to_b64(np_image: ImageType) -> str:
+        buffered = BytesIO()
+        pil_image = Image.fromarray(np.uint8(np_image[:, :, ::-1]))
+        pil_image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    def convert_b64_to_np(self, image: str) -> ImageType:
+        """
+        Converting an image as b64 encoded string into np.array
+
+        :param image: b64 encoded string
+        :return: np.array
+        """
+        return getattr(self, self.pkg_func_dict["convert_b64_to_np"])(image)
+
+    @staticmethod
+    def _cv2_convert_b64_to_np(image: str) -> ImageType:
+        np_array = np.fromstring(base64.b64decode(image), np.uint8)  # type: ignore
+        np_array = cv2.imdecode(np_array, cv2.IMREAD_COLOR).astype(np.float32)
+        return np_array.astype(uint8)
+
+    @staticmethod
+    def _pillow_convert_b64_to_np(image: str) -> ImageType:
+        array = base64.b64decode(image)
+        im_file = BytesIO(array)
+        pil_image = Image.open(im_file)
+        return np.array(pil_image)[:, :, ::-1]
+
+    def resize(self, image: ImageType, width: int, height: int, interpolation: str) -> ImageType:
+        """
+        Resize a given image to new width, height. Specifying an interpolation method is required. Depending on the
+        chosen image library use one of the following:
+
+        PIL: NEAREST, BOX, BILINEAR, BICUBIC, VIZ (available for CV2 as well)
+        CV2: INTER_NEAREST, INTER_LINEAR, INTER_AREA, VIZ
+
+        :param image: image as np.array
+        :param width: the new image width
+        :param height: the new image height
+        :param interpolation: interpolation method as string.
+        :return: resized image as np.array
+        """
+        return getattr(self, self.pkg_func_dict["resize"])(image, width, height, interpolation)
+
+    @staticmethod
+    def _cv2_resize(image: ImageType, width: int, height: int, interpolation: str) -> ImageType:
+        intpol_method_dict = {
+            "INTER_NEAREST": cv2.INTER_NEAREST,
+            "INTER_LINEAR": cv2.INTER_LINEAR,
+            "INTER_AREA": cv2.INTER_AREA,
+            "VIZ": cv2.INTER_LINEAR,
+        }
+        return cv2.resize(image, (width, height), interpolation=intpol_method_dict[interpolation])
+
+    @staticmethod
+    def _pillow_resize(image: ImageType, width: int, height: int, interpolation: str) -> ImageType:
+        intpol_method_dict = {
+            "NEAREST": Image.Resampling.NEAREST,
+            "BOX": Image.Resampling.BOX,
+            "BILINEAR": Image.Resampling.BILINEAR,
+            "BICUBIC": Image.Resampling.BICUBIC,
+            "VIZ": Image.Resampling.BILINEAR,
+        }
+        pil_image = Image.fromarray(np.uint8(image[:, :, ::-1]))
+        pil_image_resized = pil_image.resize(
+            size=(width, height), resample=intpol_method_dict[interpolation], box=None, reducing_gap=None
+        )
+        return np.array(pil_image_resized)[:, :, ::-1]
+
+    def get_text_size(self, text: str, font_scale: float) -> Tuple[int, int]:
+        """
+        Return the text size for a given font scale
+        :param text: text as string
+        :param font_scale: scale
+        :return: A tuple with width and height of the text
+        """
+        return getattr(self, self.pkg_func_dict["get_text_size"])(text, font_scale)
+
+    def _cv2_get_text_size(self, text: str, font_scale: float) -> Tuple[int, int]:
+        ((width, height), _) = cv2.getTextSize(text, self.font, font_scale, 1)  # type: ignore
+        return width, height
+
+    def _pillow_get_text_size(self, text: str, font_scale: float) -> Tuple[int, int]:  # pylint: disable=W0613
+        _, _, width, height = self.font.getbbox(text)  # type: ignore
+        return width, height
+
+    def draw_rectangle(
+        self, np_image: ImageType, box: Tuple[Any, Any, Any, Any], color: Tuple[int, int, int], thickness: int
+    ) -> ImageType:
+        """
+        Drawing a rectangle into an image with a given color (b,g,r) and given thickness
+
+        :param np_image: image
+        :param box: box (x_min, y_min, x_max, y_max)
+        :param color: (b,g,r) between 0 and 255
+        :param thickness: pixel width of the rectangle lines
+        :return: image with rectangle
+        """
+        return getattr(self, self.pkg_func_dict["draw_rectangle"])(np_image, box, color, thickness)
+
+    @staticmethod
+    def _cv2_draw_rectangle(
+        np_image: ImageType, box: Tuple[Any, Any, Any, Any], color: Sequence[int], thickness: int
+    ) -> ImageType:
+        cv2.rectangle(np_image, (box[0], box[1]), (box[2], box[3]), color=color, thickness=thickness)
+        return np_image
+
+    @staticmethod
+    def _pillow_draw_rectangle(
+        np_image: ImageType, box: Tuple[Any, Any, Any, Any], color: Sequence[int], thickness: int
+    ) -> ImageType:
+        pil_image = Image.fromarray(np.uint8(np_image[:, :, ::-1]))
+        draw = ImageDraw.Draw(pil_image)
+        draw.rectangle(box, outline=color, width=thickness)  # type: ignore
+        np_image = np.array(pil_image)[:, :, ::-1]
+        return np_image
+
+    def draw_text(
+        self, np_image: ImageType, pos: Tuple[Any, Any], text: str, color: Tuple[int, int, int], font_scale: float
+    ) -> ImageType:
+        """
+        Drawing a text into a numpy image. The result will differ between PIL and CV2 (and will not look that good when
+        using PIL).
+
+        :param np_image: image
+        :param pos: x_min, y_min position of the starting point of the text
+        :param text: text string
+        :param color: (b,g,r) between 0 and 255
+        :param font_scale: scale of font. This will only be used within a OPenCV framework
+        :return: image with text
+        """
+        return getattr(self, self.pkg_func_dict["draw_text"])(np_image, pos, text, color, font_scale)
+
+    def _cv2_draw_text(
+        self, np_image: ImageType, pos: Tuple[Any, Any], text: str, color: Tuple[int, int, int], font_scale: float
+    ) -> ImageType:
+        """
+        Draw text on an image.
+
+        :param np_image: image as np.ndarray
+        :param pos: x, y; the position of the text
+        :param text: text string to draw
+        :param color: a 3-tuple BGR color in [0, 255]
+        :param font_scale: float
+        :return: numpy array
+        """
+
+        np_image = np_image.astype(np.uint8)
+        x_0, y_0 = int(pos[0]), int(pos[1])
+        # Compute text size.
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        text_w, text_h = viz_handler.get_text_size(text, font_scale)
+        # Place text background.
+        if x_0 + text_w > np_image.shape[1]:
+            x_0 = np_image.shape[1] - text_w
+        if y_0 - int(1.15 * text_h) < 0:
+            y_0 = int(1.15 * text_h)
+        back_top_left = x_0, y_0 - int(1.3 * text_h)
+        back_bottom_right = x_0 + text_w, y_0
+        np_image = self.draw_rectangle(
+            np_image, (back_top_left[0], back_top_left[1], back_bottom_right[0], back_bottom_right[1]), color, 1
+        )
+        # Show text.
+        text_bottomleft = x_0, y_0 - int(0.25 * text_h)
+        cv2.putText(np_image, text, text_bottomleft, font, font_scale, (0, 0, 0), thickness=1, lineType=cv2.LINE_AA)
+        return np_image
+
+    @staticmethod
+    def _pillow_draw_text(
+        np_image: ImageType,
+        pos: Tuple[Any, Any],
+        text: str,
+        color: Tuple[int, int, int],  # pylint: disable=W0613
+        font_scale: float,  # pylint: disable=W0613
+    ) -> ImageType:
+        """Draw a text in an image using PIL."""
+        # using PIL default font size that does not scale to larger image sizes.
+        # Compare with https://github.com/python-pillow/Pillow/issues/6622
+        pil_image = Image.fromarray(np.uint8(np_image[:, :, ::-1]))
+        draw = ImageDraw.Draw(pil_image)
+        draw.text(pos, text, fill=(0, 0, 0), anchor="lb")
+        return np.array(pil_image)[:, :, ::-1]
+
+    def interactive_imshow(self, np_image: ImageType) -> None:
+        """Displaying an image in a separate window"""
+        return getattr(self, self.pkg_func_dict["interactive_imshow"])(np_image)
+
+    def _cv2_interactive_imshow(self, np_image: ImageType) -> None:
+        """
+        Display an image in a pop-up window
+
+        :param img: An image (expect BGR) to show.
+        """
+        name = "q, x: quit / s: save"
+        cv2.imshow(name, np_image)
+
         key = cv2.waitKey(-1)
-    key = chr(key & 0xFF)
-    cb_name = "key_cb_" + key
-    if cb_name in kwargs:
-        kwargs[cb_name](img)
-    elif key == "q":
-        cv2.destroyWindow(name)
-    elif key == "x":
-        sys.exit()
-    elif key == "s":
-        cv2.imwrite("out.png", img)
-    elif key in ["+", "="]:
-        img = cv2.resize(img, None, fx=1.3, fy=1.3, interpolation=cv2.INTER_CUBIC)
-        interactive_imshow(img, lclick_cb, rclick_cb, **kwargs)
-    elif key == "-":
-        img = cv2.resize(img, None, fx=0.7, fy=0.7, interpolation=cv2.INTER_CUBIC)
-        interactive_imshow(img, lclick_cb, rclick_cb, **kwargs)
+        while key >= 128:
+            key = cv2.waitKey(-1)
+        key = chr(key & 0xFF)  # type: ignore
+
+        if key == "q":
+            cv2.destroyWindow(name)
+        elif key == "x":
+            sys.exit()
+        elif key == "s":
+            cv2.imwrite("out.png", np_image)
+        elif key in ["+", "="]:
+            np_image = cv2.resize(np_image, None, fx=1.3, fy=1.3, interpolation=cv2.INTER_CUBIC)
+            self._cv2_interactive_imshow(np_image)
+        elif key == "-":
+            np_image = cv2.resize(np_image, None, fx=0.7, fy=0.7, interpolation=cv2.INTER_CUBIC)
+            self._cv2_interactive_imshow(np_image)
+
+    @staticmethod
+    def _pillow_interactive_imshow(np_image: ImageType) -> None:
+        name = "q, x: quit / s: save"
+        pil_image = Image.fromarray(np.uint8(np_image[:, :, ::-1]))
+        pil_image.show(name)
+
+
+# default image package
+os.environ["USE_PILLOW"] = "True"
+
+viz_handler = VizPackageHandler()