Safely Rotate Images Without Cropping

[deleomike]

Albumentations SafeRotate

The goal of this feature is to allow rotations without cropping out important corners which may contain masks, bounding boxes, or keypoints.

import cv2
import matplotlib.pyplot as plt
import albumentations as A

Test Bounding Box

bboxes = [[0, 0, 100, 100]]
bb_classes = ["test"]
id_to_name = {0: "test"}

original_keypoints = [
    (0, 0),
    (1000, 1000)
]

BOX_COLOR = (255, 0, 0) # Red
TEXT_COLOR = (255, 255, 255) # White


def visualize_bbox(img, bbox, class_name, color=BOX_COLOR, thickness=2):
    """Visualizes a single bounding box on the image"""
    x_min, y_min, w, h = bbox
    x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)
   
    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)
    
    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)    
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(
        img,
        text=class_name,
        org=(x_min, y_min - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35, 
        color=TEXT_COLOR, 
        lineType=cv2.LINE_AA,
    )
    return img


def visualize_bboxes(image, bboxes, category_ids, category_id_to_name):
    img = image.copy()
    for bbox, category_id in zip(bboxes, category_ids):
        class_name = category_id_to_name[category_id]
        img = visualize_bbox(img, bbox, class_name)
    plt.figure(figsize=(8, 8))
    plt.axis('off')
    plt.imshow(img)

KEYPOINT_COLOR = (0, 255, 0) # Green

def vis_keypoints(image, keypoints, color=KEYPOINT_COLOR, diameter=15):
    image = image.copy()

    for (x, y) in keypoints:
        cv2.circle(image, (int(x), int(y)), diameter, (0, 255, 0), -1)
        
    plt.figure(figsize=(8, 8))
    plt.axis('off')
    plt.imshow(image)

The Test Image

img = cv2.imread("cat.png")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

visualize_bboxes(img, bboxes, [0], id_to_name)
vis_keypoints(img, original_keypoints)

print("Image Shape: ", img.shape)

Image Shape:  (1053, 1070, 3)

Transform Pipelines

SafeRotateTransform = A.Compose([
    A.SafeRotate(limit=180, p=1, border_mode=cv2.BORDER_CONSTANT, mask_value=0)
], bbox_params=A.BboxParams(format='coco', label_fields=['category_ids'])
 , keypoint_params=A.KeypointParams(format='xy'))

RotateTransform = A.Compose([
    A.Rotate(limit=180, p=1, border_mode=cv2.BORDER_CONSTANT, mask_value=0)
], bbox_params=A.BboxParams(format='coco', label_fields=['category_ids'])
 , keypoint_params=A.KeypointParams(format='xy'))

A SafeRotated Image

data = SafeRotateTransform(image=img, keypoints=original_keypoints, bboxes=bboxes, category_ids=[0])
tfms_img = data["image"]
keypoints = data["keypoints"]
tfms_bboxes = data["bboxes"]
print(tfms_bboxes)
print(tfms_img.shape)


visualize_bboxes(tfms_img, tfms_bboxes, [0], id_to_name)
vis_keypoints(tfms_img, keypoints)

[(927.8357983908419, 392.83170651979276, 141.67786055840782, 138.81965455307977)]
(1053, 1070, 3)

Here the image is rotated and resized to fit the original size without cropping.

A Standard Rotated Image

data = RotateTransform(image=img, keypoints=original_keypoints, bboxes=bboxes, category_ids=[0])
tfms_img = data["image"]
keypoints = data["keypoints"]
tfms_bboxes = data["bboxes"]


visualize_bboxes(tfms_img, tfms_bboxes, [0], id_to_name)
vis_keypoints(tfms_img, keypoints)

Here the image is cropped out on the corners. This could be an issue for cropping out necessary bounding boxes, masks, and keypoints

Summary

This works very similarly to the Rotate class, except it pads the original image with a calculated value before the image is rotated.

Currently, it only supports cv2.BORDER_CONSTANT for its border mode. I can work on this for a fix in a future pull request, or include it in this one if required.

Also, I've run all the tests with pytest and added SafeRotate to the test_augmentations.py file.

Fixes

fixes #642
fixes #303

[BloodAxe]

Greetings! Thanks for your PR. I believe it would be a valuable addition to Albumentations as it was also requested in past.

Would you mind adding a few changes to your change-set, in order to ensure new code meets our coding conventions?

In particular:

1. Can you rename __rotated_img_size__, __process_image__ to have more clear name reflecting what these functions do and move them to functional.py?
2. Add new transformation to README.md (Failing unit-test will guide you how to generate new table with augmentations)
3. Ensure that PR is green (No flake & black issues remaining)

Overall, great work. Looking forward to get it merged.

[Dipet]

Good work! Thanks for PR.
Please add transform into these tests:

• test_serialization.py
  • test_augmentations_serialization
  • test_augmentations_serialization_with_custom_parameters
  • test_augmentations_for_bboxes_serialization
  • test_augmentations_for_keypoints_serialization

Also please look to test_rotate_interpolation in tests_transforms I am not sure that these implementations will use INTER_NEAREST for masks.

[deleomike]

@Dipet Thank you! I added the new feature into those unit tests.

As for the test_rotate_interpolation test, I see what you mean. I just changed the implementation to work better with these tests, and SafeRotate supports border modes now. It uses one cv2 transformation rather than padding and rotating separately which gets rid of any artifacts it had before.

[deleomike]

@BloodAxe Thank you, I look forward to using this feature in Albumentations! I moved the internal functions for SafeRotate into functional.py, and I also moved the code for key points and bounding boxes into functional as well.

I updated the readme, and It seems the PR is green. It's one commit behind the master, but I don't know if that matters. Let me know if there is anything else that should be done before the PR is approved.