-
Notifications
You must be signed in to change notification settings - Fork 0
/
createMask.py
62 lines (53 loc) · 1.82 KB
/
createMask.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
from typing import Literal, Tuple, Union
import cv2
import numpy as np
from . import Ellipse, RotatedRectangle
DRAWS = {
"rectangle": lambda image, rect, colour, thickness: cv2.drawContours(
image,
[np.int0(cv2.boxPoints(tuple(rect)))],
-1,
color=colour,
thickness=thickness,
),
"ellipse": lambda image, ellipse, colour, thickness: cv2.ellipse(
image, tuple(ellipse), color=colour, thickness=thickness
),
}
def main(
_: np.ndarray,
image: np.ndarray,
shape: Literal["ellipse", "rectangle"],
x: int,
y: int,
width: int,
height: int,
angle: float,
colour: Tuple[int, int, int],
line_thickness=2,
return_image_mode=1, # controls what image to return 0=colour image with shape overlay; 1=shape on black background; 2=pass on the input image; 3=gray image with shape overlay
) -> Tuple[object, Union[Ellipse, RotatedRectangle]]:
draw = DRAWS[shape]
shape_data = ((x, y), (width, height), angle)
if return_image_mode == 0:
# colour image with shape overlay
if len(image.shape) < 3:
ret_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
else:
ret_image = image
draw(ret_image, shape_data, colour, line_thickness)
elif return_image_mode == 1:
# shape on black background
ret_image = np.zeros_like(image)
draw(ret_image, shape_data, colour, line_thickness)
# 2 had been used to return the original image
elif return_image_mode == 3:
if len(image.shape) >= 3:
ret_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
ret_image = image
draw(ret_image, shape_data, colour, line_thickness)
else:
# pass on the input image by default
ret_image = image
return ret_image, shape_data