Repository: opencv/opencv
File: samples/dnn/alpha_matting.py (Python)

Code before:
    net = cv.dnn.readNetFromONNX(args.model, engine)
    net.setPreferableBackend(get_backend_id(args.backend))
    net.setPreferableTarget(get_target_id(args.target))
    return net


def apply_modnet(model, image):
    out = model.forward()
    alpha_mask, composite = postprocess_output(image, out)
    t, _ = model.getPerfProfile()
    label = "Inference time: %.2f ms" % (t * 1000.0 / cv.getTickFrequency())

    def draw_label(img, text, color):
        h, w = img.shape[:2]
        font_scale = max(h, w) / 1000.0
        thickness = 1
        text_size, _ = cv.getTextSize(
            text, cv.FONT_HERSHEY_SIMPLEX, font_scale, thickness
        )
        x = 10
        y = text_size[1] + 10
        cv.putText(
            img, text, (x, y), cv.FONT_HERSHEY_SIMPLEX, font_scale, color, thickness
        )

    draw_label(image, label, (0, 255, 0))
    draw_label(alpha_mask, label, (255, 255, 255))
    draw_label(composite, label, (0, 255, 0))
    cv.imshow("Input", image)
    cv.imshow("Alpha Mask", alpha_mask)
    cv.imshow("Composite", composite)


def main(func_args=None):
    args = get_args_parser(func_args)
    engine = cv.dnn.ENGINE_AUTO
    if args.backend != "default" or args.target != "cpu":
        engine = cv.dnn.ENGINE_CLASSIC

    image = cv.imread(cv.samples.findFile(args.input))
    if image is None:
        print("Failed to load the input image")
        exit(-1)

    cv.namedWindow("Input", cv.WINDOW_AUTOSIZE)
    cv.namedWindow("Alpha Mask", cv.WINDOW_AUTOSIZE)
    cv.namedWindow("Composite", cv.WINDOW_AUTOSIZE)
    cv.moveWindow("Alpha Mask", 200, 50)
    cv.moveWindow("Composite", 400, 50)

    args.model = findModel(args.model, args.sha1)

Review comment (suggestion):
Move the overlay and imshow functions into main. apply_modnet should return alpha_mask and composite

Code after:
    net = cv.dnn.readNetFromONNX(args.model, engine)
    net.setPreferableBackend(get_backend_id(args.backend))
    net.setPreferableTarget(get_target_id(args.target))
    return net


def draw_label(img, text, color):
    h, w = img.shape[:2]
    font_scale = max(h, w) / 1000.0
    thickness = 1
    text_size, _ = cv.getTextSize(text, cv.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
    x = 10
    y = text_size[1] + 10
    cv.putText(img, text, (x, y), cv.FONT_HERSHEY_SIMPLEX, font_scale, color, thickness)


def apply_modnet(args, model, image):
    inp = cv.dnn.blobFromImage(
        image, args.scale, (args.width, args.height), args.mean, swapRB=args.rgb
    )
    model.setInput(inp)
    out = model.forward()
    alpha_mask = postprocess_output(image, out)
    alpha_3ch = cv.merge([alpha_mask / 255.0, alpha_mask / 255.0, alpha_mask / 255.0])
    composite = (image.astype(np.float32) * alpha_3ch).astype(np.uint8)
    return alpha_mask, composite


def main(func_args=None):
    args = get_args_parser(func_args)
    engine = cv.dnn.ENGINE_AUTO
    if args.backend != "default" or args.target != "cpu":
        engine = cv.dnn.ENGINE_CLASSIC

    image = cv.imread(cv.samples.findFile(args.input))
    if image is None:
        print("Failed to load the input image")
        exit(-1)

    cv.namedWindow("Input", cv.WINDOW_AUTOSIZE)
    cv.namedWindow("Alpha Mask", cv.WINDOW_AUTOSIZE)
    cv.namedWindow("Composite", cv.WINDOW_AUTOSIZE)
    cv.moveWindow("Alpha Mask", 200, 50)
    cv.moveWindow("Composite", 400, 50)

    args.model = findModel(args.model, args.sha1)
    net = loadModel(args, engine)

    alpha_mask, composite = apply_modnet(args, net, image)

    t, _ = net.getPerfProfile()