Clean up, remove some transforms

app.py

@@ -1,30 +1,42 @@
-import cv2
-import inspect
-import numpy as np
 import albumentations as A
-import gradio as gr
-from typing import get_type_hints
-from PIL import Image, ImageDraw
 import base64
+import cv2
+import gradio as gr
+import inspect
 import io
-from PIL import Image
-from functools import wraps
+import numpy as np
+
 from copy import deepcopy
+from functools import wraps
+from PIL import Image, ImageDraw
+from typing import get_type_hints
+
+from utils import is_not_supported_transform
 
 
+HEADER = f"""
+<div align="center">
+    <p>
+        <img src="https://avatars.githubusercontent.com/u/57894582?s=200&v=4" alt="A" width="50" height="50" style="display:inline;">
+        <span style="font-size: 30px; vertical-align: bottom;"> lbumentations Demo ({A.__version__})</span>
+    </p>
+</div>
+"""
+
 DEFAULT_TRANSFORM = "Rotate"
 
 DEFAULT_IMAGE = "images/doctor.webp"
 DEFAULT_IMAGE_HEIGHT = 400
 DEFAULT_IMAGE_WIDTH = 600
-DEFAULT_BOXES = [[265, 121, 326, 177], [192, 169, 401, 395]]
-DEFAULT_KEYPOINTS = [
-    [(x_min + x_max) // 2, (y_min + y_max) // 2]
-    for x_min, y_min, x_max, y_max in DEFAULT_BOXES
+DEFAULT_BOXES = [
+    [265, 121, 326, 177],  # Mask
+    [192, 169, 401, 395],  # Coverall
 ]
-CORENERS = [[[x_min, y_min], [x_max, y_max], [x_min, y_max], [x_max, y_min]] for x_min, y_min, x_max, y_max in DEFAULT_BOXES]
-for bbox_corners in CORENERS:
-    DEFAULT_KEYPOINTS += bbox_corners
+
+mask_keypoints = [[270, 123], [320, 130], [270, 151], [321, 158]]
+pocket_keypoints = [[226, 379], [272, 386], [307, 388], [364, 380]]
+arm_keypoints = [[215, 194], [372, 192], [214, 322], [378, 330]]
+DEFAULT_KEYPOINTS = mask_keypoints + pocket_keypoints + arm_keypoints
 
 BASE64_DEFAULT_MASKS = [
     {
@@ -45,15 +57,23 @@ BASE64_DEFAULT_MASKS = [
 transforms_map = {
     name: cls
     for name, cls in vars(A).items()
-    if inspect.isclass(cls) and issubclass(cls, (A.DualTransform, A.ImageOnlyTransform))
+    if (
+        inspect.isclass(cls)
+        and issubclass(cls, (A.DualTransform, A.ImageOnlyTransform))
+        and not is_not_supported_transform(cls)
+    )
 }
 transforms_map.pop("DualTransform", None)
 transforms_map.pop("ImageOnlyTransform", None)
+transforms_map.pop("ReferenceBasedTransform", None)
 transforms_keys = list(sorted(transforms_map.keys()))
 
+
 # Decode the masks
 for mask in BASE64_DEFAULT_MASKS:
-    mask["mask"] = np.array(Image.open(io.BytesIO(base64.b64decode(mask["mask"]))).convert("L"))
+    mask["mask"] = np.array(
+        Image.open(io.BytesIO(base64.b64decode(mask["mask"]))).convert("L")
+    )
 
 
 def run_with_retry(compose):
@@ -80,6 +100,7 @@ def run_with_retry(compose):
                 raise e
         compose.processors = processors
         return result
+
     return wrapper
 
 
@@ -118,12 +139,12 @@ def draw_mask(image, mask):
     return image_with_mask
 
 
-def draw_not_implemented_image(image):
+def draw_not_implemented_image(image: np.ndarray, annotation_type: str):
     """Draw the image with a text. In the middle."""
     pil_image = Image.fromarray(image)
     draw = ImageDraw.Draw(pil_image)
     # align in the centerm, and make bigger font
-    text = "NOT IMPLEMETED FOR THIS TYPE OF ANNOTATIONS"
+    text = f'Transform NOT working with "{annotation_type.upper()}" annotaions.'
     length = draw.textlength(text)
     draw.text(
         (DEFAULT_IMAGE_WIDTH // 2 - length // 2, DEFAULT_IMAGE_HEIGHT // 2),
@@ -164,45 +185,65 @@ def get_formatted_signature(function_or_class, indentation=4):
     return result
 
 
-def update(image, code):
-    try:
-        augmentation = eval(code)
-        compose = A.Compose(
-            [augmentation],
-            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_id"]),
-            keypoint_params=A.KeypointParams(format="xy"),
-            additional_targets={"not_implemented_image": "image"}
-        )
-        compose = run_with_retry(compose)  # to prevent NotImplementedError
-
-        keypoints = DEFAULT_KEYPOINTS
-        bboxes = DEFAULT_BOXES
-        mask = get_rgb_mask(BASE64_DEFAULT_MASKS)
-        augmented = compose(
-            image=image,
-            not_implemented_image=draw_not_implemented_image(image),
-            mask=mask,
-            keypoints=keypoints,
-            bboxes=bboxes,
-            category_id=range(len(bboxes)),
-        )
-        image = augmented["image"]
-        not_implemented_image = augmented["not_implemented_image"]
-        mask = augmented.get("mask", None)
-        bboxes = augmented.get("bboxes", None)
-        keypoints = augmented.get("keypoints", None)
+def get_formatted_transform(transform_number):
+    transform_name = transforms_keys[transform_number]
+    transform = transforms_map[transform_name]
+    return f"A.{transform.__name__}{get_formatted_signature(transform)}"
+
+
+def get_formatted_transform_docs(transform_number):
+    transform_name = transforms_keys[transform_number]
+    transform = transforms_map[transform_name]
+    return transform.__doc__.strip("\n")
 
-        image_with_mask = draw_mask(image.copy(), mask) if mask is not None else not_implemented_image
-        image_with_bboxes = draw_boxes(image.copy(), bboxes) if bboxes is not None else not_implemented_image
-        image_with_keypoints = draw_keypoints(image.copy(), keypoints) if keypoints is not None else not_implemented_image
 
-        return [
-            (image_with_mask, "Mask"),
-            (image_with_bboxes, "Boxes"),
-            (image_with_keypoints, "Keypoints"),
-        ]
-    except Exception as e:
-        raise e
+def update_augmented_images(image, code):
+
+    augmentation = eval(code)
+    compose = A.Compose(
+        [augmentation],
+        bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_id"]),
+        keypoint_params=A.KeypointParams(format="xy"),
+    )
+    compose = run_with_retry(compose)  # to prevent NotImplementedError
+
+    keypoints = DEFAULT_KEYPOINTS
+    bboxes = DEFAULT_BOXES
+    mask = get_rgb_mask(BASE64_DEFAULT_MASKS)
+    augmented = compose(
+        image=image,
+        not_implemented_image=image.copy(),
+        mask=mask,
+        keypoints=keypoints,
+        bboxes=bboxes,
+        category_id=range(len(bboxes)),
+    )
+    image = augmented["image"]
+    mask = augmented.get("mask", None)
+    bboxes = augmented.get("bboxes", None)
+    keypoints = augmented.get("keypoints", None)
+
+    # Draw the augmented images (or replace by placeholder if not implemented)
+    if mask is not None:
+        image_with_mask = draw_mask(image.copy(), mask)
+    else:
+        image_with_mask = draw_not_implemented_image(image.copy(), "mask")
+
+    if bboxes is not None:
+        image_with_bboxes = draw_boxes(image.copy(), bboxes)
+    else:
+        image_with_bboxes = draw_not_implemented_image(image.copy(), "boxes")
+
+    if keypoints is not None:
+        image_with_keypoints = draw_keypoints(image.copy(), keypoints)
+    else:
+        image_with_keypoints = draw_not_implemented_image(image.copy(), "keypoints")
+
+    return [
+        (image_with_mask, "Mask"),
+        (image_with_bboxes, "Boxes"),
+        (image_with_keypoints, "Keypoints"),
+    ]
 
 
 def update_image_info(image):
@@ -212,20 +253,14 @@ def update_image_info(image):
     return f"Image info:\n\t - shape: {h}x{w}\n\t - dtype: {dtype}\n\t - min/max: {min_}/{max_}"
 
 
-def get_formatted_transform(transform_number):
-    transform_name = transforms_keys[transform_number]
-    transform = transforms_map[transform_name]
-    return f"A.{transform.__name__}{get_formatted_signature(transform)}"
-
-
-def get_formatted_transform_docs(transform_number):
-    transform_name = transforms_keys[transform_number]
-    transform = transforms_map[transform_name]
-    return transform.__doc__.strip("\n")
+def update_code_and_docs(select):
+    code = get_formatted_transform(select)
+    docs = get_formatted_transform_docs(select)
+    return code, docs
 
 
 with gr.Blocks() as demo:
-
+    gr.Markdown(HEADER)
     with gr.Row():
         with gr.Column():
             with gr.Group():
@@ -236,7 +271,7 @@ with gr.Blocks() as demo:
                     type="index",
                     interactive=True,
                 )
-                with gr.Accordion("Documentation", open=False):
+                with gr.Accordion("Documentation (click to expand)", open=False):
                     docs = gr.TextArea(
                         get_formatted_transform_docs(
                             transforms_keys.index(DEFAULT_TRANSFORM)
@@ -245,8 +280,11 @@ with gr.Blocks() as demo:
                         interactive=False,
                     )
                 code = gr.Code(
+                    label="Code",
                     language="python",
-                    value=get_formatted_transform(transforms_keys.index(DEFAULT_TRANSFORM)),
+                    value=get_formatted_transform(
+                        transforms_keys.index(DEFAULT_TRANSFORM)
+                    ),
                     interactive=True,
                     lines=5,
                 )
@@ -256,7 +294,7 @@ with gr.Blocks() as demo:
                 lines=1,
                 max_lines=1,
             )
-            button = gr.Button("Run")
+            button = gr.Button("Apply!")
         image = gr.Image(
             value=DEFAULT_IMAGE,
             type="numpy",
@@ -266,11 +304,11 @@ with gr.Blocks() as demo:
         )
     with gr.Row():
         augmented_image = gr.Gallery(rows=1, columns=3)
-        # augmented_image = gr.Image(type="numpy", height=300, width=300)
 
-    #image.upload(fn=update_image_info, inputs=[image], outputs=[info])
-    select.change(fn=get_formatted_transform, inputs=[select], outputs=[code])
-    button.click(fn=update, inputs=[image, code], outputs=[augmented_image])
+    select.change(fn=update_code_and_docs, inputs=[select], outputs=[code, docs])
+    button.click(
+        fn=update_augmented_images, inputs=[image, code], outputs=[augmented_image]
+    )
 
 
 if __name__ == "__main__":

utils.py

@@ -0,0 +1,31 @@
+import albumentations as A
+import inspect
+from typing import Callable
+
+
+FILTER_TRANSFORMS = [
+    A.ImageOnlyTransform,
+    A.DualTransform,
+    A.ReferenceBasedTransform,
+    A.TemplateTransform,
+    A.Lambda,
+]
+
+
+def is_not_supported_transform(transform_cls):
+    sig = inspect.signature(transform_cls)
+
+    if issubclass(transform_cls, A.ReferenceBasedTransform):
+        return True
+
+    for filter_transform_cls in FILTER_TRANSFORMS:
+        if transform_cls is filter_transform_cls:
+            return True
+
+    for param in sig.parameters.values():
+        if issubclass(type(param.annotation), type(Callable)):
+            return True
+        if param.name in ["read_fn", "reference_images"]:
+            return True
+
+    return False