Repository: kornia/kornia

Issue:
[Feat] Add more draw utilities
## 🚀 Feature
Implement more drawing functions in `kornia.utils`:
- [ ] `draw_point`
- [ ] `draw_line`
- [ ] `draw_polyline` https://github.com/kornia/kornia/issues/278

Some resources:
- https://stackoverflow.com/questions/31638651/how-can-i-draw-lines-into-numpy-arrays



Discussion:
Can I help with this one?
I've been dabbling with `draw_line` for a couple of days, and I finished writing the function for it using this [Line Drawing Algorithm](https://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm#Line_equation), at the time it is not optimized and not properly validated.

I've created a [kaggle notebook](https://www.kaggle.com/abdelrhmanhosny/test-draw-line)  to show how my function compares to opencv's function of draw line.

If this implementation is accepted, I can modify it to have proper validations, test cases and optimize it.
I think the goal was to do that on GPU. I am not sure whether numpy suits or under optimization you mean moving that to GPU?
@Abdelrhman-Hosny the implementation looks really cool and seems pretty straightforward to port it to pure torch operators and will cover already @dobosevych comments.

I think the ideal would be something like this:
```python
def draw_line(image: torch.Tensor, line: torch.Tensor) -> torch.Tensor:
"""Draw a tensor containing defined geometric 2d lines into an image.

  Args:
    image: the input image to where to draw the lines with shape (B,C,H,W).
    lines: the tensor containing the lines to draw with shape (B, N, 4) where lines are defined [x1 y1 x2 y2].

  Return:
     the image containing the lines.
"""
```
We can think also about adding other properties like colour, thickness, etc.

Any other thoughts @shijianjian @ducha-aiki @anguelos  ?
One more comment - we have `render_gaussian2d` that is meant to render 2d points in a differentiable way, but it's quite memory consuming and not really sure that's a good choice here. @anibali  thoughts ?
https://kornia.readthedocs.io/en/latest/geometry.subpix.html#kornia.geometry.subpix.render_gaussian2d
@edgarriba @dobosevych I just tried it, and the function can take both numpy and pytorch arrays, so it should work on GPUs. (Added the change to the [kaggle notebook](https://www.kaggle.com/abdelrhmanhosny/test-draw-line) in case you need to check it out).

What I meant by optimization was to use cython and to get rid of the floating point operations, which should make this run faster.

Adding color is easy, as for thickness, I have an idea that I tried, but didn't thoroughly test it.

Should I start working on it? (turning it into the above shape)
and if yes 
- is the cython part needed?
- The function above takes, B images and for each image there are N lines, is that correct ?
>     * is the cython part needed?
kornia only accepts torch as dependency

>     * The function above takes, B images and for each image there are N lines, is that correct ?
yeap :+1: 
I've started working, now I've added a color parameter to the line.

I've created 3 functions

```
def draw_pixel(
    image: torch.Tensor,
    x: int,
    y: int,
    color: torch.Tensor,
    ):
    r"""Draws a pixel into an image.

    Args:
        image: the input image to where to draw the lines with shape (C,H,W).
        x: the x coordinate of the pixel.
        y: the y coordinate of the pixel.
        color: the color of the pixel with shape (3).
    
    Return:
        Nothing is returned
"""
```
****

```
def draw_lines(
    images: torch.Tensor,
    lines: torch.Tensor,
    color: torch.Tensor,
    ) -> torch.Tensor:
    r"""Draw a tensor containing defined geometric 2d lines into an image.

    Args:
        images: the input image to where to draw the lines with shape (B,C,H,W).
        lines: the tensor containing the lines to draw with shape (B, N, 4) where lines are defined [x1 y1 x2 y2].

    Return:
        the images (B,C,H,W) containing the lines.
    """
```
****

```
def draw_line(
    image : torch.Tensor,
    p1 : torch.Tensor, p2 : torch.Tensor,
    color : torch.Tensor,
    ):
    r"""Draw a single line into an image.

    Args:
        image: the input image to where to draw the lines with shape (C,H,W).
        p1: the start point of the line with shape (2).
        p2: the end point of the line with shape (2).
        color: the color of the line with shape (3).

    Return:
        the image containing the line.
    """
```
****

I'm writing to give an update, also confirm the names and inputs of the functions, are these okay ?.

I would also like to ask whether all 3 functions should be added to the `kornia.utils` or whether I should only add `draw_lines` and the rest be put somewhere else ?

I tested these manually and got the results I wanted, I'm waiting for confirmation so I could write automatic tests.
I think having the three functions public it's good. One question - `draw_pixel` why returns `None`, is it supposed to happen inplace ? why not return the image ?
My thought process was that `draw_pixel` only has one line that is executed `image[:, y, x] = color` and the rest of it is just will just be checks to make sure that the pixels are inside the image. 

So, it won't be used outside another function. should I change it to `_draw_pixel` instead ?
In that case, yes - keep it private

Resolving patch:
diff --git a/kornia/utils/__init__.py b/kornia/utils/__init__.py
--- a/kornia/utils/__init__.py
+++ b/kornia/utils/__init__.py
@@ -1,4 +1,4 @@
-from .draw import draw_rectangle
+from .draw import draw_line, draw_rectangle
 from .grid import create_meshgrid, create_meshgrid3d
 from .helpers import _extract_device_dtype, safe_inverse_with_mask, safe_solve_with_mask
 from .image import image_list_to_tensor, image_to_tensor, ImageToTensor, tensor_to_image
@@ -18,6 +18,7 @@
     "save_pointcloud_ply",
     "load_pointcloud_ply",
     "draw_rectangle",
+    "draw_line",
     "_extract_device_dtype",
     "safe_inverse_with_mask",
     "safe_solve_with_mask",
diff --git a/kornia/utils/draw.py b/kornia/utils/draw.py
--- a/kornia/utils/draw.py
+++ b/kornia/utils/draw.py
@@ -5,6 +5,134 @@
 # TODO: implement width of the line
 
 
+def _draw_pixel(
+    image: torch.Tensor,
+    x: int,
+    y: int,
+    color: torch.Tensor,
+) -> None:
+    r"""Draws a pixel into an image.
+
+    Args:
+        image: the input image to where to draw the lines with shape :math`(C,H,W)`.
+        x: the x coordinate of the pixel.
+        y: the y coordinate of the pixel.
+        color: the color of the pixel with :math`(C)` where :math`C` is the number of channels of the image.
+
+    Return:
+        Nothing is returned.
+    """
+    image[:, y, x] = color
+
+
+def draw_line(
+    image: torch.Tensor,
+    p1: torch.Tensor,
+    p2: torch.Tensor,
+    color: torch.Tensor,
+) -> torch.Tensor:
+    r"""Draw a single line into an image.
+
+    Args:
+        image: the input image to where to draw the lines with shape :math`(C,H,W)`.
+        p1: the start point [x y] of the line with shape (2).
+        p2: the end point [x y] of the line with shape (2).
+        color: the color of the line with shape :math`(C)` where :math`C` is the number of channels of the image.
+
+    Return:
+        the image with containing the line.
+
+    Examples:
+    >>> image = torch.zeros(1, 8, 8)
+    >>> draw_line(image, torch.tensor([6, 4]), torch.tensor([1, 4]), torch.tensor([255]))
+    tensor([[[  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
+             [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
+             [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
+             [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
+             [  0., 255., 255., 255., 255., 255., 255.,   0.],
+             [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
+             [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
+             [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.]]])
+    """
+
+    if (len(p1) != 2) or (len(p2) != 2):
+        raise ValueError("p1 and p2 must have length 2.")
+
+    if len(image.size()) != 3:
+        raise ValueError("image must have 3 dimensions (C,H,W).")
+
+    if (color.size(0) != image.size(0)):
+        raise ValueError("color must have the same number of channels as the image.")
+
+    if (p1[0] >= image.size(2)) or (p1[1] >= image.size(1) or (p1[0] < 0) or (p1[1] < 0)):
+        raise ValueError("p1 is out of bounds.")
+
+    if (p2[0] >= image.size(2)) or (p2[1] >= image.size(1) or (p2[0] < 0) or (p2[1] < 0)):
+        raise ValueError("p2 is out of bounds.")
+
+    # assign points
+    x1, y1 = p1
+    x2, y2 = p2
+
+    # calcullate coefficients A,B,C of line
+    # from equation Ax + By + C = 0
+    A = y2 - y1
+    B = x1 - x2
+    C = x2 * y1 - x1 * y2
+
+    # make sure A is positive to utilize the functiom properly
+    if (A < 0):
+        A = -A
+        B = -B
+        C = -C
+
+    # calculate the slope of the line
+    # check for division by zero
+    if (B != 0):
+        m = -A / B
+
+    # make sure you start drawing in the right direction
+    x1, x2 = min(x1, x2), max(x1, x2)
+    y1, y2 = min(y1, y2), max(y1, y2)
+
+    # line equation that determines the distance away from the line
+    def line_equation(x, y):
+        return A * x + B * y + C
+
+    # vertical line
+    if B == 0:
+        image[:, y1:y2 + 1, x1] = color
+    # horizontal line
+    elif A == 0:
+        image[:, y1, x1:x2 + 1] = color
+    # slope between 0 and 1
+    elif 0 < m < 1:
+        for i in range(x1, x2 + 1):
+            _draw_pixel(image, i, y1, color)
+            if line_equation(i + 1, y1 + 0.5) > 0:
+                y1 += 1
+    # slope greater than or equal to 1
+    elif m >= 1:
+        for j in range(y1, y2 + 1):
+            _draw_pixel(image, x1, j, color)
+            if line_equation(x1 + 0.5, j + 1) < 0:
+                x1 += 1
+    # slope less then -1
+    elif m <= -1:
+        for j in range(y1, y2 + 1):
+            _draw_pixel(image, x2, j, color)
+            if line_equation(x2 - 0.5, j + 1) > 0:
+                x2 -= 1
+    # slope between -1 and 0
+    elif -1 < m < 0:
+        for i in range(x1, x2 + 1):
+            _draw_pixel(image, i, y2, color)
+            if line_equation(i + 1, y2 - 0.5) > 0:
+                y2 -= 1
+
+    return image
+
+
 def draw_rectangle(
     image: torch.Tensor,
     rectangle: torch.Tensor,


Test patch:
diff --git a/test/utils/test_draw.py b/test/utils/test_draw.py
--- a/test/utils/test_draw.py
+++ b/test/utils/test_draw.py
@@ -2,6 +2,157 @@
 import torch
 
 from kornia.utils import draw_rectangle
+from kornia.utils.draw import draw_line
+
+
+class TestDrawLine:
+    def test_draw_line_vertical(self, dtype, device):
+        """Test drawing a vertical line."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        img = draw_line(img, torch.tensor([6, 2]), torch.tensor([6, 0]), torch.tensor([255]))
+        img_mask = img == torch.tensor([[0., 0., 0., 0., 0., 0., 255., 0.],
+                                        [0., 0., 0., 0., 0., 0., 255., 0.],
+                                        [0., 0., 0., 0., 0., 0., 255., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.]])
+        assert torch.all(img_mask)
+
+    def test_draw_line_horizontal(self, dtype, device):
+        """Test drawing a horizontal line."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        img = draw_line(img, torch.tensor([6, 4]), torch.tensor([0, 4]), torch.tensor([255]))
+        img_mask = img == torch.tensor([[0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [255., 255., 255., 255., 255., 255., 255., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.]])
+        assert torch.all(img_mask)
+
+    def test_draw_line_m_lte_neg1(self, dtype, device):
+        """Test drawing a line with m <= -1."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        img = draw_line(img, torch.tensor([0, 7]), torch.tensor([6, 0]), torch.tensor([255]))
+        img_mask = img == torch.tensor([[0., 0., 0., 0., 0., 0., 255., 0.],
+                                        [0., 0., 0., 0., 0., 255., 0., 0.],
+                                        [0., 0., 0., 0., 255., 0., 0., 0.],
+                                        [0., 0., 0., 255., 0., 0., 0., 0.],
+                                        [0., 0., 0., 255., 0., 0., 0., 0.],
+                                        [0., 0., 255., 0., 0., 0., 0., 0.],
+                                        [0., 255., 0., 0., 0., 0., 0., 0.],
+                                        [255., 0., 0., 0., 0., 0., 0., 0.]])
+        assert torch.all(img_mask)
+
+    def test_draw_line_m_lt_0_gte_neg1(self, dtype, device):
+        """Test drawing a line with -1 < m < 0."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        img = draw_line(img, torch.tensor([1, 5]), torch.tensor([7, 0]), torch.tensor([255]))
+        img_mask = img == torch.tensor([[0., 0., 0., 0., 0., 0., 0., 255.],
+                                        [0., 0., 0., 0., 0., 0., 255., 0.],
+                                        [0., 0., 0., 0., 0., 255., 0., 0.],
+                                        [0., 0., 0., 255., 255., 0., 0., 0.],
+                                        [0., 0., 255., 0., 0., 0., 0., 0.],
+                                        [0., 255., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.]])
+        assert torch.all(img_mask)
+
+    def test_draw_line_m_gt_0_lt_1(self, dtype, device):
+        """Test drawing a line with 0 < m < 1."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        img = draw_line(img, torch.tensor([0, 0]), torch.tensor([6, 2]), torch.tensor([255]))
+        img_mask = img == torch.tensor([[255., 255., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 255., 255., 255., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 255., 255., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.]])
+        assert torch.all(img_mask)
+
+    def test_draw_line_m_gte_1(self, dtype, device):
+        """Test drawing a line with m >= 1."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        img = draw_line(img, torch.tensor([3, 7]), torch.tensor([1, 4]), torch.tensor([255]))
+        img_mask = img == torch.tensor([[0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 0., 0., 0., 0., 0.],
+                                        [0., 255., 0., 0., 0., 0., 0., 0.],
+                                        [0., 0., 255., 0., 0., 0., 0., 0.],
+                                        [0., 0., 255., 0., 0., 0., 0., 0.],
+                                        [0., 0., 0., 255., 0., 0., 0., 0.]])
+        assert torch.all(img_mask)
+
+    @pytest.mark.parametrize('p1',
+                             [
+                                 torch.tensor([-1, 0]),
+                                 torch.tensor([0, -1]),
+                                 torch.tensor([8, 0]),
+                                 torch.tensor([0, 8]),
+                             ])
+    def test_p1_out_of_bounds(self, p1, dtype, device):
+        """Tests that an exception is raised if p1 is out of bounds."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        with pytest.raises(ValueError) as excinfo:
+            draw_line(img, p1, torch.tensor([0, 0]), torch.tensor([255]))
+
+        assert 'p1 is out of bounds.' == str(excinfo.value)
+
+    @pytest.mark.parametrize('p2',
+                             [
+                                 torch.tensor([-1, 0]),
+                                 torch.tensor([0, -1]),
+                                 torch.tensor([8, 0]),
+                                 torch.tensor([0, 8]),
+                             ])
+    def test_p2_out_of_bounds(self, p2, dtype, device):
+        """Tests that an exception is raised if p2 is out of bounds."""
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        with pytest.raises(ValueError) as excinfo:
+            draw_line(img, torch.tensor([0, 0]), p2, torch.tensor([255]))
+
+        assert 'p2 is out of bounds.' == str(excinfo.value)
+
+    @pytest.mark.parametrize('img_size', [(200, 100), (32, 3, 20, 20)])
+    def test_image_size(self, img_size, dtype, device):
+        img = torch.zeros(*img_size, dtype=dtype, device=device)
+        with pytest.raises(ValueError) as excinfo:
+            draw_line(img, torch.tensor([0, 0]), torch.tensor([1, 1]), torch.tensor([255]))
+
+        assert 'image must have 3 dimensions (C,H,W).' == str(excinfo.value)
+
+    @pytest.mark.parametrize('img_size,color',
+                             [
+                                 ((1, 8, 8), torch.tensor([23, 53])),
+                                 ((3, 8, 8), torch.tensor([255]))
+                             ])
+    def test_color_image_channel_size(self, img_size, color, dtype, device):
+        img = torch.zeros(*img_size, dtype=dtype, device=device)
+        with pytest.raises(ValueError) as excinfo:
+            draw_line(img, torch.tensor([0, 0]), torch.tensor([1, 1]), color)
+
+        assert 'color must have the same number of channels as the image.' == str(excinfo.value)
+
+    @pytest.mark.parametrize('p1,p2',
+                             [
+                                 ((0, 1), (1, 5, 2)),
+                                 ((0, 1, 2), (1, 5)),
+                                 (torch.tensor([0, 1]), torch.tensor([0, 2, 3])),
+                                 (torch.tensor([0, 1, 5]), torch.tensor([0, 2])),
+                             ])
+    def test_point_size(self, p1, p2, dtype, device):
+        img = torch.zeros(1, 8, 8, dtype=dtype, device=device)
+        with pytest.raises(ValueError) as excinfo:
+            draw_line(img, p1, p2, torch.tensor([255]))
+
+        assert 'p1 and p2 must have length 2.' == str(excinfo.value)
 
 
 class TestDrawRectangle:
@@ -65,8 +216,8 @@ def test_n_rectangles(self, device, batch, N, fill):
                         random_w_rectangle[
                             b,
                             :,
-                            points_list[b][n][1] : points_list[b][n][3] + 1,
-                            points_list[b][n][0] : points_list[b][n][2] + 1,
+                            points_list[b][n][1]: points_list[b][n][3] + 1,
+                            points_list[b][n][0]: points_list[b][n][2] + 1,
                         ].sum()
                         == (points_list[b][n][3] - points_list[b][n][1] + 1)
                         * (points_list[b][n][2] - points_list[b][n][0] + 1)
@@ -75,25 +226,25 @@ def test_n_rectangles(self, device, batch, N, fill):
                 else:
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1] : points_list[b][n][3] + 1, points_list[b][n][0]
+                            b, :, points_list[b][n][1]: points_list[b][n][3] + 1, points_list[b][n][0]
                         ].sum()
                         == (points_list[b][n][3] - points_list[b][n][1] + 1) * 3
                     )
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1] : points_list[b][n][3] + 1, points_list[b][n][2]
+                            b, :, points_list[b][n][1]: points_list[b][n][3] + 1, points_list[b][n][2]
                         ].sum()
                         == (points_list[b][n][3] - points_list[b][n][1] + 1) * 3
                     )
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1], points_list[b][n][0] : points_list[b][n][2] + 1
+                            b, :, points_list[b][n][1], points_list[b][n][0]: points_list[b][n][2] + 1
                         ].sum()
                         == (points_list[b][n][2] - points_list[b][n][0] + 1) * 3
                     )
                     assert (
                         random_w_rectangle[
-                            b, :, points_list[b][n][1], points_list[b][n][0] : points_list[b][n][2] + 1
+                            b, :, points_list[b][n][1], points_list[b][n][0]: points_list[b][n][2] + 1
                         ].sum()
                         == (points_list[b][n][2] - points_list[b][n][0] + 1) * 3
                     )