Last active
June 27, 2022 23:27
-
-
Save petered/ddcf471bd65eda69c9fe8ae96113fc82 to your computer and use it in GitHub Desktop.
Simple tool using OpenCV to manually draw bounding boxes on an image, and return the coordinates of those boxes
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from typing import Optional, Tuple, Sequence | |
import cv2 | |
def draw_bboxes_on_image( | |
image, # A (height, width, 3) unit8 numpy array representing an image | |
return_relative=False, # True to return box coorinates relative to width/height of image so they will be in [0, 1] interval | |
window_name='Draw Your Bounding Box', # What to call the window | |
display_color=(0, 0, 255), # Colour of the displayed box | |
display_thickness=2 # Thickness of the displayed box | |
) -> Sequence[Tuple[float, float, float, float]]: | |
""" | |
Draw a bounding box on the image, returning a (x_left, x_right, y_top, y_bottom) coordinates. | |
Once this function is launched, an image will pop up. | |
- Click and drag to add a new bounding box | |
- Click 'Return' to close the window and return the currently shown boxes | |
- Click 'c' (for clear) to clear all boxes and start again | |
- Click 'b' (for back) to remove the last drawn box | |
- Click 'q' (for quit) to close the window and return no boxes | |
""" | |
class CallbackContainer: | |
bboxes = [] | |
start_xy: Optional[Tuple[float, float]] = None | |
@classmethod | |
def on_click(cls, event, x, y, flags, param): | |
print(f'Got event type {cv2.EVENT_LBUTTONDOWN} at {x}, {y}') | |
if event == cv2.EVENT_LBUTTONDOWN: | |
if cls.start_xy is None: | |
cls.start_xy = (x, y) | |
elif event == cv2.EVENT_LBUTTONUP: | |
xs, xe = sorted([cls.start_xy[0], x]) | |
ys, ye = sorted([cls.start_xy[1], y]) | |
cls.bboxes.append([xs, xe, ys, ye]) | |
print(f'Added box #{len(cls.bboxes)}: {cls.bboxes[-1]}') | |
cls.show_display_image() | |
cls.start_xy = None | |
@classmethod | |
def show_display_image(cls): | |
display_image = image.copy() | |
for xs, xe, ys, ye in cls.bboxes: | |
cv2.rectangle(display_image, (xs, ys), (xe, ye), color=display_color, thickness=display_thickness) | |
cv2.imshow(window_name, display_image) | |
cv2.namedWindow(window_name) | |
cv2.setMouseCallback(window_name, CallbackContainer.on_click) | |
CallbackContainer.show_display_image() | |
while True: | |
print('Click and drag to make a box on the image') | |
key = cv2.waitKey() | |
if key == ord('c'): | |
print('Clearted all boxes') | |
CallbackContainer.bboxes.clear() | |
elif key == ord('b'): | |
print('Clearing last box') | |
if len(CallbackContainer.bboxes)>0: | |
CallbackContainer.bboxes.pop(-1) | |
elif key == ord('q'): | |
print('Quitting, returning no boxes') | |
CallbackContainer.bboxes.clear() | |
break | |
elif key == 13: # Enter | |
print(f'Done, returning {len(CallbackContainer.bboxes)} boxes.') | |
break | |
else: | |
print(f'Unknown keypress: {key}') | |
CallbackContainer.show_display_image() | |
cv2.destroyWindow(window_name) | |
boxes = CallbackContainer.bboxes.copy() | |
if return_relative: | |
h, w = image.shape[:2] | |
boxes = [[xs / w, xe / w, ys / h, ye / h] for xs, xe, ys, ye in boxes] | |
print(f'Your boxes are\n {boxes}') | |
return boxes | |
if __name__ == "__main__": | |
draw_bboxes_on_image( | |
image=cv2.imread(cv2.samples.findFile('starry_night.jpg')) | |
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment