diff --git a/samples/dnn/alpha_matting.cpp b/samples/dnn/alpha_matting.cpp new file mode 100644 index 0000000000..950e26d37a --- /dev/null +++ b/samples/dnn/alpha_matting.cpp @@ -0,0 +1,201 @@ +/* + * This file is part of OpenCV project. + * It is subject to the license terms in the LICENSE file found in the top-level directory of this distribution and at http://opencv.org/license.html. + * + * Copyright (C) 2025, Bigvision LLC. + * + * @file alpha_matting.cpp + * @brief MODNet Alpha Matting using OpenCV DNN + * + * This sample demonstrates human portrait alpha matting using MODNet model. + * MODNet is a trimap-free portrait matting method that can produce high-quality + * alpha mattes for portrait images in real-time. + * + * Reference: + * Github: https://github.com/ZHKKKe/MODNet + * + * Usage: + * ./example_dnn_alpha_matting --input=image.jpg # Process image + * + * Requirements: + * - OpenCV >= 5.0.0 with DNN module + * - MODNet ONNX model + */ + +#include +#include +#include +#include +#include +#include + +#include "common.hpp" + +using namespace cv; +using namespace cv::dnn; +using namespace std; + +const string about = + "This sample demonstrates human portrait alpha matting using MODNet model.\n" + "MODNet is a trimap-free portrait matting method that can produce high-quality\n" + "alpha mattes for portrait images in real-time.\n\n" + "Usage examples:\n" + "\t./example_alpha_matting --input=image.jpg\n" + "\t./example_alpha_matting modnet (using config alias)\n\n" + "To download the MODNet model, run: python download_models.py modnet\n" + "Press any key to exit \n"; + + +const string param_keys = + "{ help h | | Print help message }" + "{ @alias | modnet | An alias name of model to extract preprocessing parameters from models.yml file }" + "{ zoo | ../dnn/models.yml | An optional path to file with preprocessing parameters }" + "{ input i | messi5.jpg | Path to input image file }" + "{ model | | Path to MODNet ONNX model file }"; + +const string backend_keys = format( + "{ backend | default | Choose one of computation backends: " + "default: automatically (by default), " + "openvino: Intel's Deep Learning Inference Engine, " + "opencv: OpenCV implementation, " + "vkcom: VKCOM, " + "cuda: CUDA, " + "webnn: WebNN }"); + +const string target_keys = format( + "{ target | cpu | Choose one of target computation devices: " + "cpu: CPU target (by default), " + "opencl: OpenCL, " + "opencl_fp16: OpenCL fp16 (half-float precision), " + "vpu: VPU, " + "vulkan: Vulkan, " + "cuda: CUDA, " + "cuda_fp16: CUDA fp16 (half-float precision) }"); + +string keys = param_keys + backend_keys + target_keys; + +static void loadModel(const string modelPath, String backend, String target, Net &net, EngineType engine) +{ + net = readNetFromONNX(modelPath, engine); + net.setPreferableBackend(getBackendID(backend)); + net.setPreferableTarget(getTargetID(target)); +} + +static void postprocess(const Mat &image, const Mat &alpha_output, Mat &alpha_mask) +{ + int h = image.rows; + int w = image.cols; + + Mat alpha; + if (alpha_output.dims == 4 && alpha_output.size[0] == 1 && alpha_output.size[1] == 1) + { + alpha = alpha_output.reshape(0, {alpha_output.size[2], alpha_output.size[3]}); + } + else + { + alpha = alpha_output.clone(); + } + + resize(alpha, alpha, Size(w, h)); + + alpha = cv::min(cv::max(alpha, 0.0), 1.0); + alpha.convertTo(alpha_mask, CV_8U, 255.0); +} + +static void processImage(const Mat &image, Mat &alpha_mask, Mat &composite, Net &net, + float scale, int width, int height, const Scalar &mean, bool swapRB) +{ + if (image.empty()) + return; + + Mat blob = blobFromImage(image, scale, Size(width, height), mean, swapRB, false, CV_32F); + net.setInput(blob); + Mat output = net.forward(); + postprocess(image, output, alpha_mask); + + Mat alpha_3ch; + cvtColor(alpha_mask, alpha_3ch, COLOR_GRAY2BGR); + alpha_3ch.convertTo(alpha_3ch, CV_32F, 1.0 / 255.0); + + Mat image_f; + image.convertTo(image_f, CV_32F); + multiply(image_f, alpha_3ch, composite); + composite.convertTo(composite, CV_8U); +} + +static void setupWindows() +{ + namedWindow("Original", WINDOW_AUTOSIZE); + namedWindow("Alpha Mask", WINDOW_AUTOSIZE); + namedWindow("Composite", WINDOW_AUTOSIZE); + moveWindow("Alpha Mask", 200, 0); + moveWindow("Composite", 400, 0); +} + +int main(int argc, char **argv) +{ + CommandLineParser parser(argc, argv, keys); + + if (parser.has("help")) + { + cout << about << endl; + parser.printMessage(); + return 0; + } + + string modelName = parser.get("@alias"); + string zooFile = parser.get("zoo"); + + zooFile = findFile(zooFile); + + keys += genPreprocArguments(modelName, zooFile); + + parser = CommandLineParser(argc, argv, keys); + + int input_width = parser.get("width"); + int input_height = parser.get("height"); + float scale_factor = parser.get("scale"); + Scalar mean_values = parser.get("mean"); + bool swapRB = parser.get("rgb"); + String backend = parser.get("backend"); + String target = parser.get("target"); + String sha1 = parser.get("sha1"); + + string model = findModel(parser.get("model"), sha1); + + parser.about(about); + + EngineType engine = ENGINE_AUTO; + if (backend != "default" || target != "cpu") + { + engine = ENGINE_CLASSIC; + } + + Net net; + loadModel(model, backend, target, net, engine); + + string input_path = samples::findFile(parser.get("input")); + Mat image = imread(input_path); + if (image.empty()) + { + cout << "[ERROR] Cannot load input image: " << input_path << endl; + return -1; + } + + setupWindows(); + + cout << "Processing image: " << input_path << endl; + cout << "Press any key to exit" << endl; + + Mat alpha_mask, composite; + + processImage(image, alpha_mask, composite, net, scale_factor, input_width, input_height, mean_values, swapRB); + + imshow("Original", image); + imshow("Alpha Mask", alpha_mask); + imshow("Composite", composite); + + waitKey(0); + destroyAllWindows(); + return 0; +} diff --git a/samples/dnn/alpha_matting.py b/samples/dnn/alpha_matting.py new file mode 100644 index 0000000000..0c31f1055e --- /dev/null +++ b/samples/dnn/alpha_matting.py @@ -0,0 +1,177 @@ +""" +This file is part of OpenCV project. +It is subject to the license terms in the LICENSE file found in the top-level directory of this distribution and at http://opencv.org/license.html. + +Copyright (C) 2025, Bigvision LLC. + +MODNet Alpha Matting with OpenCV DNN + +This sample demonstrates human portrait alpha matting using MODNet model. +MODNet is a trimap-free portrait matting method that can produce high-quality +alpha mattes for portrait images in real-time. + +Reference: + Github: https://github.com/ZHKKKe/MODNet + +To download the MODNet model, run: + python download_models.py modnet + +Usage: + python alpha_matting.py --input=image.jpg +""" + +import cv2 as cv +import numpy as np +import argparse +import os +from common import * + + +def get_args_parser(func_args): + backends = ("default", "openvino", "opencv", "vkcom", "cuda") + targets = ( + "cpu", + "opencl", + "opencl_fp16", + "ncs2_vpu", + "hddl_vpu", + "vulkan", + "cuda", + "cuda_fp16", + ) + + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument( + "--zoo", + default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "models.yml"), + help="An optional path to file with preprocessing parameters.", + ) + parser.add_argument( + "--input", + default="messi5.jpg", + help="Path to input image or video file. Defaults to messi5.jpg in samples/data.", + ) + parser.add_argument( + "--backend", + default="default", + type=str, + choices=backends, + help="Choose one of computation backends: " + "default: automatically (by default), " + "openvino: Intel's Deep Learning Inference Engine, " + "opencv: OpenCV implementation, " + "vkcom: VKCOM, " + "cuda: CUDA", + ) + parser.add_argument( + "--target", + default="cpu", + type=str, + choices=targets, + help="Choose one of target computation devices: " + "cpu: CPU target (by default), " + "opencl: OpenCL, " + "opencl_fp16: OpenCL fp16 (half-float precision), " + "ncs2_vpu: NCS2 VPU, " + "hddl_vpu: HDDL VPU, " + "vulkan: Vulkan, " + "cuda: CUDA, " + "cuda_fp16: CUDA fp16 (half-float precision)", + ) + + args, _ = parser.parse_known_args() + add_preproc_args(args.zoo, parser, "alpha_matting", "modnet") + parser = argparse.ArgumentParser( + parents=[parser], + description=""" + To run: + python alpha_matting.py --input=path/to/your/input/image + + Model path can also be specified using --model argument + """, + formatter_class=argparse.RawTextHelpFormatter, + ) + return parser.parse_args(func_args) + + +def postprocess_output(image, alpha_output): + """Process model output to create alpha mask.""" + h, w = image.shape[:2] + + alpha = alpha_output[0, 0] if alpha_output.ndim == 4 else alpha_output[0] + alpha = cv.resize(alpha, (w, h)) + alpha = np.clip(alpha, 0, 1) + + alpha_mask = (alpha * 255).astype(np.uint8) + + return alpha_mask + + +def loadModel(args, engine): + 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() + label = "Inference time: %.2f ms" % (t * 1000.0 / cv.getTickFrequency()) + + 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) + + print("Press any key to exit") + cv.waitKey(0) + cv.destroyAllWindows() + + +if __name__ == "__main__": + main() diff --git a/samples/dnn/models.yml b/samples/dnn/models.yml index b2d7f59e75..8ffc92cb46 100644 --- a/samples/dnn/models.yml +++ b/samples/dnn/models.yml @@ -506,3 +506,18 @@ NAFNet: scale: 0.00392 rgb: true sample: "deblurring" + +################################################################################ +# Alpha Matting model. +################################################################################ +modnet: + load_info: + url: "https://github.com/HarxSan/Modnet/raw/main/modnet.onnx" + sha1: "40eebf4387ea86c982bf6e363a7f84b659b145a0" + model: "modnet.onnx" + mean: [0.0, 0.0, 0.0] + scale: 0.00784313725 + width: 512 + height: 512 + rgb: true + sample: "alpha_matting"