demo.py

import argparse
import numpy as np
import cv2 as cv
from efficientSAM import EfficientSAM

# Check OpenCV version
opencv_python_version = lambda str_version: tuple(map(int, (str_version.split("."))))
assert opencv_python_version(cv.__version__) >= opencv_python_version("4.10.0"), \
       "Please install latest opencv-python for benchmark: python3 -m pip install --upgrade opencv-python"

# Valid combinations of backends and targets
backend_target_pairs = [
    [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_TARGET_CPU],
    [cv.dnn.DNN_BACKEND_CUDA,   cv.dnn.DNN_TARGET_CUDA],
    [cv.dnn.DNN_BACKEND_CUDA,   cv.dnn.DNN_TARGET_CUDA_FP16],
    [cv.dnn.DNN_BACKEND_TIMVX,  cv.dnn.DNN_TARGET_NPU],
    [cv.dnn.DNN_BACKEND_CANN,   cv.dnn.DNN_TARGET_NPU]
]

parser = argparse.ArgumentParser(description='EfficientSAM Demo')
parser.add_argument('--input', '-i', type=str,
                    help='Set input path to a certain image.')
parser.add_argument('--model', '-m', type=str, default='image_segmentation_efficientsam_ti_2025april.onnx',
                    help='Set model path, defaults to image_segmentation_efficientsam_ti_2025april.onnx.')
parser.add_argument('--backend_target', '-bt', type=int, default=0,
                    help='''Choose one of the backend-target pair to run this demo:
                        {:d}: (default) OpenCV implementation + CPU,
                        {:d}: CUDA + GPU (CUDA),
                        {:d}: CUDA + GPU (CUDA FP16),
                        {:d}: TIM-VX + NPU,
                        {:d}: CANN + NPU
                    '''.format(*[x for x in range(len(backend_target_pairs))]))
parser.add_argument('--save', '-s', action='store_true',
                    help='Specify to save a file with results. Invalid in case of camera input.')
args = parser.parse_args()

# Global configuration
WINDOW_SIZE = (800, 600)  # Fixed window size (width, height)
MAX_POINTS = 6             # Maximum allowed points
points = []                # Store clicked coordinates (original image scale)
labels = []                # Point labels (-1: useless, 0: background, 1: foreground, 2: top-left, 3: bottom right)
backend_point = []
rectangle = False
current_img = None

def visualize(image, result):
    """
    Visualize the inference result on the input image.

    Args:
        image (np.ndarray): The input image.
        result (np.ndarray): The inference result.

    Returns:
        vis_result (np.ndarray): The visualized result.
    """
    # get image and mask
    vis_result = np.copy(image)
    mask = np.copy(result)
    # change mask to binary image
    t, binary = cv.threshold(mask, 127, 255, cv.THRESH_BINARY)
    assert set(np.unique(binary)) <= {0, 255}, "The mask must be a binary image."
    # enhance red channel to make the segmentation more obviously
    enhancement_factor = 1.8
    red_channel = vis_result[:, :, 2]
    # update the channel
    red_channel = np.where(binary == 255, np.minimum(red_channel * enhancement_factor, 255), red_channel)
    vis_result[:, :, 2] = red_channel

    # draw borders
    contours, hierarchy = cv.findContours(binary, cv.RETR_LIST, cv.CHAIN_APPROX_TC89_L1)
    cv.drawContours(vis_result, contours, contourIdx = -1, color = (255,255,255), thickness=2)
    return vis_result

def select(event, x, y, flags, param):
    """Handle mouse events with coordinate conversion"""
    global points, labels, backend_point, rectangle, current_img
    orig_img = param['original_img']
    image_window = param['image_window']

    if event == cv.EVENT_LBUTTONDOWN:
        param['mouse_down_time'] = cv.getTickCount()
        backend_point = [x, y]

    elif event == cv.EVENT_MOUSEMOVE:
        if rectangle == True:
            rectangle_change_img = current_img.copy()
            cv.rectangle(rectangle_change_img, (backend_point[0], backend_point[1]), (x, y), (255,0,0) , 2)
            cv.imshow(image_window, rectangle_change_img)
        elif len(backend_point) != 0 and len(points) < MAX_POINTS:
            rectangle = True


    elif event == cv.EVENT_LBUTTONUP:
        if len(points) >= MAX_POINTS:
            print(f"Maximum points reached {MAX_POINTS}.")
            return

        if rectangle == False:
            duration = (cv.getTickCount() - param['mouse_down_time'])/cv.getTickFrequency()
            label = -1 if duration > 0.5 else 1  # Long press = background

            points.append([backend_point[0], backend_point[1]])
            labels.append(label)
            print(f"Added {['background','foreground','background'][label]} point {backend_point}.")
        else:
            if len(points) + 1 >= MAX_POINTS:
                rectangle = False
                backend_point.clear()
                cv.imshow(image_window, current_img)
                print(f"Points reached {MAX_POINTS}, could not add box.")
                return
            point_leftup = []
            point_rightdown = []
            if x > backend_point[0] or y > backend_point[1]:
                point_leftup.extend(backend_point)
                point_rightdown.extend([x,y])
            else:
                point_leftup.extend([x,y])
                point_rightdown.extend(backend_point)
            points.append(point_leftup)
            points.append(point_rightdown)
            print(f"Added box from {point_leftup} to {point_rightdown}.")
            labels.append(2)
            labels.append(3)
            rectangle = False
        backend_point.clear()

        marked_img = orig_img.copy()
        top_left = None
        for (px, py), lbl in zip(points, labels):
            if lbl == -1:
                cv.circle(marked_img, (px, py), 5, (0, 0, 255), -1)
            elif lbl == 1:
                cv.circle(marked_img, (px, py), 5, (0, 255, 0), -1)
            elif lbl == 2:
                top_left = (px, py)
            elif lbl == 3:
                bottom_right = (px, py)
                cv.rectangle(marked_img, top_left, bottom_right, (255,0,0) , 2)
        cv.imshow(image_window, marked_img)
        current_img = marked_img.copy()


if __name__ == '__main__':
    backend_id = backend_target_pairs[args.backend_target][0]
    target_id = backend_target_pairs[args.backend_target][1]
    # Load the EfficientSAM model
    model = EfficientSAM(modelPath=args.model)

    if args.input is not None:
        # Read image
        image = cv.imread(args.input)
        if image is None:
            print('Could not open or find the image:', args.input)
            exit(0)
        # create window
        image_window = "Origin image"
        cv.namedWindow(image_window, cv.WINDOW_NORMAL)
        # change window size
        rate = 1
        rate1 = 1
        rate2 = 1
        if(image.shape[1]>WINDOW_SIZE[0]):
            rate1 = WINDOW_SIZE[0]/image.shape[1]
        if(image.shape[0]>WINDOW_SIZE[1]):
            rate2 = WINDOW_SIZE[1]/image.shape[0]
        rate = min(rate1, rate2)
        # width, height
        WINDOW_SIZE = (int(image.shape[1] * rate), int(image.shape[0] * rate))
        cv.resizeWindow(image_window, WINDOW_SIZE[0], WINDOW_SIZE[1])
        # put the window on the left of the screen
        cv.moveWindow(image_window, 50, 100)
        # set listener to record user's click point
        param = {
            'original_img': image,
            'mouse_down_time': 0,
            'image_window' : image_window
        }
        cv.setMouseCallback(image_window, select, param)
        # tips in the terminal
        print("Click — Select foreground point\n"
        "Long press — Select background point\n"
        "Drag — Create selection box\n"
        "Enter — Infer\n"
        "Backspace — Clear the prompts\n"
        "Q - Quit")
        # show image
        cv.imshow(image_window, image)
        current_img = image.copy()
        # create window to show visualized result
        vis_image = image.copy()
        segmentation_window = "Segment result"
        cv.namedWindow(segmentation_window, cv.WINDOW_NORMAL)
        cv.resizeWindow(segmentation_window, WINDOW_SIZE[0], WINDOW_SIZE[1])
        cv.moveWindow(segmentation_window, WINDOW_SIZE[0]+51, 100)
        cv.imshow(segmentation_window, vis_image)
        # waiting for click
        while True:
            # Check window status
            # if click × to close the image window then ending
            if (cv.getWindowProperty(image_window, cv.WND_PROP_VISIBLE) < 1 or
                cv.getWindowProperty(segmentation_window, cv.WND_PROP_VISIBLE) < 1):
                break

            # Handle keyboard input
            key = cv.waitKey(1)

            # receive enter
            if key == 13:

                vis_image = image.copy()
                cv.putText(vis_image, "infering...",
                            (50, vis_image.shape[0]//2),
                            cv.FONT_HERSHEY_SIMPLEX, 10, (255,255,255), 5)
                cv.imshow(segmentation_window, vis_image)

                result = model.infer(image=image, points=points, labels=labels)
                if len(result) == 0:
                    print("clear and select points again!")
                else:
                    vis_result = visualize(image, result)

                    cv.imshow(segmentation_window, vis_result)
            elif key == 8 or key == 127:  # ASCII for Backspace or Delete
                points.clear()
                labels.clear()
                backend_point = []
                rectangle = False
                current_img = image
                print("Points are cleared.")
                cv.imshow(image_window, image)
            elif key == ord('q') or key == ord('Q'):
                break

        cv.destroyAllWindows()

        # Save results if save is true
        if args.save:
            cv.imwrite('./example_outputs/vis_result.jpg', vis_result)
            cv.imwrite("./example_outputs/mask.jpg", result)
            print('vis_result.jpg and mask.jpg are saved to ./example_outputs/')

    else:
        print('Set input path to a certain image.')
        pass