fix: Fixed rotation of img + target

[fg-mindee]

Following the refactoring of rotation, some edge cases weren't handled. This PR restarts fresh and introduces the following modifications:

• deleted rotate to introduce rotate_sample
• the function can take boxes or polygons, with or without expansion
• unittests were considerably extended to avoid further issues
• added box filtering in RandomRotate

Here is a snippet to show this:

from doctr.transforms.functional import rotate_sample
import matplotlib.pyplot as plt
import numpy as np
import torch

# Get the inputs
img = torch.ones((3, 200, 100), dtype=torch.float32)
h, w = img.shape[1:]
box = np.array([0, 0, .5, .25])[None, ...]
target = np.zeros((*img.shape[1:], 3), dtype=np.uint8)
for xmin, ymin, xmax, ymax in box:
    target[int(ymin * h): int(ymax * h), int(xmin * w): int(xmax * w)] = (255, 255, 255)

# Rotate without expansion
out_img, out_poly = rotate_sample(img, box, 90, expand=False)
h, w = out_img.shape[1:]
# Make an image for the target
out_target = np.zeros((*out_img.shape[1:], 3), dtype=np.uint8)
out_box = np.concatenate((out_poly.min(1), out_poly.max(1)), -1)
for xmin, ymin, xmax, ymax in out_box:
    out_target[int(ymin * h): int(ymax * h), int(xmin * w): int(xmax * w)] = (255, 255, 255)

# Rotate with expansion
exp_img, exp_poly = rotate_sample(img, box, 90, expand=True)
h, w = exp_img.shape[1:]
exp_target = np.zeros((*exp_img.shape[1:], 3), dtype=np.uint8)
exp_box = np.concatenate((exp_poly.min(1), exp_poly.max(1)), -1)
for xmin, ymin, xmax, ymax in exp_box:
    exp_target[int(ymin * h): int(ymax * h), int(xmin * w): int(xmax * w)] = (255, 255, 255)

_, axes = plt.subplots(2, 3)
# First col = Inputs
axes[0, 0].imshow(img.permute(1, 2, 0).numpy())
axes[1, 0].imshow(target)
# 2nd col = no expansion
axes[0, 1].imshow(out_img.permute(1, 2, 0).numpy())
axes[1, 1].imshow(out_target)
# 3rd col = expansion
axes[0, 2].imshow(exp_img.permute(1, 2, 0).numpy())
axes[1, 2].imshow(exp_target)

plt.show()

which produces:

top row = image
bot row = target
1st col = original
2nd col = 90° rotation without expansion
3rd col = 90° rotation with expansion

(without expansion, the box is out of the image)

Any feedback is welcome!

[codecov]

Codecov Report

Merging #784 (df5235f) into main (8d271ea) will increase coverage by 0.05%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #784      +/-   ##
==========================================
+ Coverage   96.00%   96.05%   +0.05%     
==========================================
  Files         130      130              
  Lines        4882     4898      +16     
==========================================
+ Hits         4687     4705      +18     
+ Misses        195      193       -2     

Flag	Coverage Δ
unittests	96.05% <100.00%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
doctr/transforms/functional/pytorch.py	100.00% <100.00%> (ø)
doctr/transforms/functional/tensorflow.py	100.00% <100.00%> (+2.27%)	⬆️
doctr/transforms/modules/base.py	94.59% <100.00%> (+0.07%)	⬆️
doctr/utils/geometry.py	98.83% <100.00%> (+1.01%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8d271ea...df5235f. Read the comment docs.

[charlesmindee]

Thanks for the PR!