Initial

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+*.pyo
+
+.venv/

app.py

@@ -0,0 +1,252 @@
+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 io
+from PIL import Image
+from functools import wraps
+
+
+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
+]
+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
+
+BASE64_DEFAULT_MASKS = [
+    {
+        "label": "Coverall",
+        # light green color
+        "color": (144, 238, 144),
+        "mask": "iVBORw0KGgoAAAANSUhEUgAAAlgAAAGQCAAAAABXXkFEAAAF+ElEQVR4nO3dwXLjNhBFUSg1///LziLj1Iwt26KkFhuvz9kkWVhFAJcAqXGStQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACe5+3sCyCUsq745+wLSKCsz4T1DMr6RFiUENZT2LI+EhYlhPWnt7t3nvt/MtSvsy+gkcfaeFuXJ11HBJPx7r+s7piPP3o0m/9zFP729tdfjv/gnT8dy1G41npeEc7Dd3astR7q6uOP2rT+4wb7mMLBGfkckildyyw8HMa1HWr8pK7pz1hF55YnrdE315dVHZmTb9IcPLVr7Oi/36oOTMpPe97Q+R16FL7wzW3sqThw2DdkdfOs3JTowDmeN+gbN6tbp+XJHxdk1pAPnIE3TczNnzdrmteaNeJjj1Y3zMyRD5w00WuNGu/RR/afpubZn5dlymjvexH8Znbu+cApk73WlLE+8v3C1Rm69wNnTPdaM0ba6hcOJkz4WhPG2SqrtSZM+Vr5o2yX1Vr5k75W+hhbZrVW+rSvlT3Ctlmt7Hlfa0X/anLnrnpf3DPkhtV86Zpf3sNyw+ouvKzYsPqvW/8rfERsWJxLWOeJ3rKERQlhnSh5y0oNK3nNtpAa1h6C8w8NK3jFNpEZlq5OlxnWNnLvgMiwcpdrH5FhbST2HkgMK3axdpIY1lZS7wJhUSIwrM32gM0u91aBYdFBXlihO8Bu4sLSVQ9xYe0n81ZICytzlTaUFtaOIm8GYVEiLKzIm39LYWHRhbAoIawGEg9wYVEiK6zEW39TWWHtKvCGEBYlhEUJYbWQdxYKixLCokRUWHkHyr6iwqIPYVEiKSwnYSNJYdGIsCghLEoIixLCokRQWF4KOwkKa2txd4WwmkgrS1iUEFYXYVuWsCiRE1bYHb+7nLBoRViUEBYlhEUJYVFCWG1kvdYKixIxYWXd7/tLCUtXzaSEFeBy9gU8VUhYNqxuMsLSVTsZYdFORFgZG1bGKN5FhEU/wqJEQlhZZ0iIhLBoSFiUEBYlhEUJYVEiIaysP70NkRAWDQmLEsLqI+qLXmFRQliUEBYlhEWJX2dfwK4ua4U9bj+XsA66/P0P0vqCo/CQy8dv+X3r/wVhHXElI2VdJ6wDiiOKalRYlBDWo6L2mecRVhtZhQrrUb5wuEpYlBBWF1knYUZYYWsSISKsM8vyiHVdRlivYWM8ICQsa95NSFinleUk/EJKWDQjrCbSDvOYsM5ZGCfhV2LCohdhPcKG9SVh3e5TRk/sKu0RKyis1y+N/eobOWG9nK6+ExTWa7esN119y79XeBdV/URYd5DVz4R1wNtlqepGUa+5+6551DKstaIe3ulEWJQQFiWERQlhUSIqrLx3q31FhUUfwqJEVljOwjaywqINYXUQuNOGhRW4QpsKC4su0sKyZTWRFtaWZe14zT+JCytylTYUuQyb/cJf5Brk7Vhrt5Xa62pvFRnWVmu107UekBlW6mptJDQsZZ0tNaxtpN4BqeN6fzW8/PH3LaUuQOq4PuqaVuz8OwopEXvHfNRyywqe/eCh/aVjV9Fz7z8KcpborIR1jvCo1hLWGQZk5a3wBCO6EtbLzehKWNQQFiWERQlhUUJYlBAWJYRFCWG9Wsc/Di8gLEoMCWvINtHIkLB4NWFRQliUENbLzXjeExYlhEUJYVFiRlgzHmtamREWLycsSgjr9UYczMKihLAoISxKCIsSwqKEsCghLEoI6wQTvsgSFiWERQlhUUJYZxjwkCUsSgiLEiPCGnDytDMirH7yU58QVv4qNjQhLE4gLEoIixLCooSwzhH/QjEhrMuQ/31NKxPCktYJBs14s9MnfOZn7FhrLdvWaw0KS1qvNG+qu5yI4TMfPryrmqSVPfWjjsLfnIgvMHWOO+xa0XM/ccdaK3xRO5gaFsWERQlhUWJqWB0e3qNNDauD6LiFRYmhYUVvFi3MDEtX5UaGpat6I8Oi3r8KSpCuwVpGmQAAAABJRU5ErkJggg==",
+    },
+    {
+        "label": "Mask",
+        # light blue color
+        "color": (173, 216, 230),
+        "mask": "iVBORw0KGgoAAAANSUhEUgAAAlgAAAGQCAAAAABXXkFEAAAB4ElEQVR4nO3csQ6CMBSG0avv/864OFhoobW9UeM5i4ML+fOFkmCMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIbcPn0BX257ftppkMEqtuY35uplqYN2VhFhsU5m2rnIKsJmXYy00xFWRBjuin1KvV2F6c7dP30Bv2ugwT8krPcp64SwJmzSahJWYbQUZbUIqzD8QK6sBmEVdLKKsEghLFIIa5LDs05YpBAWKYQ1y1lYJSxSCIsUwiKFsF55XlpGWLP83q9KWKQQ1qt37j6OzyphkUJYpBBWwZP4KsIqKWsRYZFCWDtuWWsIixTC2nPLWkJYB8paQVhHY2XpsMosdV0vaozXZpsG/+s3x0BN9bQM1sdOJ45pmauXpU4VadlqgLEubRF2AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAGCZBwKLGEVAl/J/AAAAAElFTkSuQmCC",
+    },
+]
+
+# Get all the transforms from the albumentations library
+transforms_map = {
+    name: cls
+    for name, cls in vars(A).items()
+    if inspect.isclass(cls) and issubclass(cls, (A.DualTransform, A.ImageOnlyTransform))
+}
+transforms_map.pop("DualTransform", None)
+transforms_map.pop("ImageOnlyTransform", 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"))
+
+
+def run_with_retry(compose):
+    @wraps(compose)
+    def wrapper(*args, **kwargs):
+        for i in range(4):
+            try:
+                return compose(*args, **kwargs)
+            except NotImplementedError as e:
+                print(f"Caught NotImplementedError: {e}")
+                if "bbox" in str(e):
+                    kwargs.pop("bboxes", None)
+                    kwargs.pop("category_id", None)
+                if "keypoint" in str(e):
+                    kwargs.pop("keypoints", None)
+                if "mask" in str(e):
+                    kwargs.pop("mask", None)
+    return wrapper
+
+
+
+def draw_boxes(image, boxes, color=(255, 0, 0), thickness=2) -> np.ndarray:
+    """Draw boxes with PIL."""
+    pil_image = Image.fromarray(image)
+    draw = ImageDraw.Draw(pil_image)
+    for box in boxes:
+        x_min, y_min, x_max, y_max = box
+        draw.rectangle([x_min, y_min, x_max, y_max], outline=color, width=thickness)
+    return np.array(pil_image)
+
+
+def draw_keypoints(image, keypoints, color=(255, 0, 0), radius=2):
+    """Draw keypoints with PIL."""
+    pil_image = Image.fromarray(image)
+    draw = ImageDraw.Draw(pil_image)
+    for keypoint in keypoints:
+        x, y = keypoint
+        draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=color)
+    return np.array(pil_image)
+
+
+def get_rgb_mask(masks):
+    """Get the RGB mask from the binary mask."""
+    rgb_mask = np.zeros((DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH, 3), dtype=np.uint8)
+    for data in masks:
+        mask = data["mask"]
+        rgb_mask[mask > 0] = np.array(data["color"])
+    return rgb_mask
+
+
+def draw_mask(image, mask):
+    """Draw the mask on the image."""
+    image_with_mask = cv2.addWeighted(image, 0.5, mask, 0.5, 0)
+    return image_with_mask
+
+
+def draw_not_implemented_image(image):
+    """Draw the image with a text. In the middle."""
+    pil_image = Image.fromarray(image)
+    draw = ImageDraw.Draw(pil_image)
+    draw.text((DEFAULT_IMAGE_WIDTH // 2, DEFAULT_IMAGE_HEIGHT // 2), "Not implemented", fill=(255, 0, 0))
+    return np.array(pil_image)
+
+
+def get_formatted_signature(function_or_class, indentation=4):
+
+    signature = inspect.signature(function_or_class)
+    type_hints = get_type_hints(function_or_class)
+
+    args = []
+    for param in signature.parameters.values():
+        if param.default == inspect.Parameter.empty:
+            str_param = f"{param.name}=,"
+        else:
+            if isinstance(param.default, str):
+                str_param = f'{param.name}="{param.default}",'
+            else:
+                str_param = f"{param.name}={param.default},"
+
+        annotation = type_hints.get(param.name, param.annotation)
+        if isinstance(param.annotation, type):
+            str_param += f"  # {param.annotation.__name__}"
+        else:
+            str_annotation = str(annotation).replace("typing.", "")
+            str_param += f"  # {str_annotation}"
+        str_param = "\n" + " " * indentation + str_param
+        args.append(str_param)
+
+    result = "(" + "".join(args) + "\n" + " " * (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)
+
+        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_image_info(image):
+    h, w = image.shape[:2]
+    dtype = image.dtype
+    max_, min_ = image.max(), image.min()
+    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")
+
+
+with gr.Blocks() as demo:
+
+    with gr.Row():
+        with gr.Group():
+            select = gr.Dropdown(
+                label="Select a transformation",
+                choices=transforms_keys,
+                value=DEFAULT_TRANSFORM,
+                type="index",
+                interactive=True,
+            )
+            with gr.Accordion("Documentation", open=False):
+                docs = gr.TextArea(
+                    get_formatted_transform_docs(
+                        transforms_keys.index(DEFAULT_TRANSFORM)
+                    ),
+                    show_label=False,
+                    interactive=False,
+                )
+            code = gr.Code(
+                language="python",
+                value=get_formatted_transform(transforms_keys.index(DEFAULT_TRANSFORM)),
+                interactive=True,
+                lines=5,
+            )
+            #info = gr.Text(interactive=False, label="Image info", value="")
+        image = gr.Image(
+            value=DEFAULT_IMAGE,
+            type="numpy",
+            height=500,
+            width=300,
+            sources=[],
+        )
+    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 = gr.Button("Run")
+    button.click(fn=update, inputs=[image, code], outputs=[augmented_image])
+
+
+if __name__ == "__main__":
+    demo.launch()

packages.txt

@@ -0,0 +1 @@
+python3-opencv

requirements.txt

@@ -0,0 +1,2 @@
+albumentations
+Pillow