Annotator for points and boxes on the image for SAM using cv2 (keyboard shortcuts: p = point selection mode | b = box selection mode | q = quit)

sam_annotator.py

def annotate_image(image):
    # Create a copy of the image to draw annotations on
    annotated_image = image.copy()

    # Create lists to store the annotated points and boxes
    points = []
    boxes = []
    point_labels = []
    unscaled_points = []  # To store unscaled points
    unscaled_boxes = []  # To store unscaled boxes

    # Flag to indicate annotation mode (True for points, False for boxes)
    point_mode = True

    # Zoom parameters
    zoom_factor = 1.0
    zoom_scale = 0.1

    # Mouse callback function
    def annotate_callback(event, x, y, flags, param):
        nonlocal points, boxes, point_labels, unscaled_points, unscaled_boxes, annotated_image, point_mode, zoom_factor

        if event == cv2.EVENT_LBUTTONDOWN:
            # Add clicked point to the list if in point annotation mode
            if point_mode:
                unscaled_points.append((x, y))
                points.append((int(x / zoom_factor), int(y / zoom_factor)))  # Scale the point and store
                # Draw the annotated point on the image
                cv2.circle(annotated_image, (x, y), 5, (0, 255, 0), -1)
                cv2.imshow('Annotated Image', cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))

                point_labels.append(1)  # Add the label 1 to the point_labels list

            # Start recording box annotation if in box annotation mode
            else:
                points.append((x, y))
                unscaled_points.append((int(x / zoom_factor), int(y / zoom_factor)))  # Scale the point and store

        elif event == cv2.EVENT_RBUTTONDOWN:
            # Add clicked point to the list with label 0
            unscaled_points.append((x, y))
            points.append((int(x / zoom_factor), int(y / zoom_factor)))  # Scale the point and store
            # Draw the annotated point on the image
            cv2.circle(annotated_image, (x, y), 5, (255, 0, 0), -1)
            cv2.imshow('Annotated Image', cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))

            point_labels.append(0)  # Add the label 0 to the point_labels list

        elif event == cv2.EVENT_LBUTTONUP:
            # Add box coordinates to the list if in box annotation mode
            if not point_mode:
                xmin = min(points[-1][0], x)
                ymin = min(points[-1][1], y)
                xmax = max(points[-1][0], x)
                ymax = max(points[-1][1], y)
                unscaled_boxes.append([xmin, ymin, xmax, ymax])  # Store unscaled box
                boxes.append(
                    [int(coord / zoom_factor) for coord in [xmin, ymin, xmax, ymax]])  # Scale the box and store

                # Draw the annotated box on the image
                cv2.rectangle(annotated_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
                cv2.imshow('Annotated Image', cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))

                # Remove the last point from the points if it was used to draw a box
                points.pop()
                unscaled_points.pop()

        elif event == cv2.EVENT_MOUSEWHEEL:
            # Zoom in or out based on the mouse wheel movement
            if flags > 0:
                zoom_factor += zoom_scale
            else:
                zoom_factor -= zoom_scale

            # Apply the zoom transformation to the image and update the displayed image
            zoomed_image = cv2.resize(image, None, fx=zoom_factor, fy=zoom_factor, interpolation=cv2.INTER_LINEAR)
            annotated_image = zoomed_image.copy()

            # Adjust the coordinates of the points and boxes based on the zoom factor
            adjusted_points = [(int(x * zoom_factor), int(y * zoom_factor)) for x, y in points]
            adjusted_boxes = [[int(coord * zoom_factor) for coord in box] for box in boxes]

            # Redraw the points and boxes on the zoomed image
            for point in adjusted_points:
                cv2.circle(annotated_image, point, 5, (0, 255, 0), -1)
            for box in adjusted_boxes:
                cv2.rectangle(annotated_image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
            cv2.imshow('Annotated Image', cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))

    # Create a named window and set the mouse callback
    cv2.namedWindow('Annotated Image')
    cv2.setMouseCallback('Annotated Image', annotate_callback)

    # Display the initial image
    cv2.imshow('Annotated Image', cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    while True:
        # Wait for a key press
        key = cv2.waitKey(1) & 0xFF

        # Switch to box annotation mode if 'b' is pressed
        if key == ord('b'):
            point_mode = False
            print('Switched to box annotation mode')

        # Switch to point annotation mode if 'p' is pressed
        elif key == ord('p'):
            point_mode = True
            print('Switched to point annotation mode')

        # Quit if 'q' is pressed
        elif key == ord('q'):
            break

    # Close all windows
    cv2.destroyAllWindows()

    # Return the annotated points and boxes
    return points, boxes, point_labels


# Usage example
# right click - annotate a point with label 0 (background)
# left click - annotate a point with label 1 (foreground)
# left click + drag - annotate a box
# scroll wheel - zoom in or out

# keyboard shortcuts:
# 'b' - switch to box annotation mode
# 'p' - switch to point annotation mode
# 'q' - quit