mirror of
https://github.com/opencv/opencv.git
synced 2025-08-05 22:19:14 +08:00
Merge pull request #26736 from gursimarsingh:inpainting_onnx_model
Added lama inpainting onnx model sample #26736 ### Pull Request Readiness Checklist See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request - [x] I agree to contribute to the project under Apache 2 License. - [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV - [x] The PR is proposed to the proper branch - [x] There is a reference to the original bug report and related work - [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Gursimar Singh
GitHub
parent
55a2ca58f0
commit
a27b90c217
230
samples/dnn/inpainting.cpp
Normal file
230
samples/dnn/inpainting.cpp
Normal file
@ -0,0 +1,230 @@
|
||||
/*
|
||||
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.
|
||||
|
||||
This sample inpaints the masked area in the given image.
|
||||
|
||||
Copyright (C) 2025, Bigvision LLC.
|
||||
|
||||
How to use:
|
||||
Sample command to run:
|
||||
|
||||
./example_dnn_inpainting
|
||||
The system will ask you to draw the mask on area to be inpainted
|
||||
|
||||
You can download lama inpainting model using:
|
||||
`python download_models.py lama`
|
||||
|
||||
References:
|
||||
Github: https://github.com/advimman/lama
|
||||
ONNX model: https://huggingface.co/Carve/LaMa-ONNX/blob/main/lama_fp32.onnx
|
||||
|
||||
ONNX model was further quantized using block quantization from [opencv_zoo](https://github.com/opencv/opencv_zoo)
|
||||
|
||||
Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to point to the directory where models are downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/dnn.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace dnn;
|
||||
using namespace std;
|
||||
|
||||
const string about = "Use this script for image inpainting using OpenCV. \n\n"
|
||||
"Firstly, download required models i.e. lama using `download_models.py` (if not already done). Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to point to the directory where models are downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.\n"
|
||||
"To run:\n"
|
||||
"\t Example: ./example_dnn_inpainting [--input=<image_name>] \n\n"
|
||||
"Inpainting model path can also be specified using --model argument.\n\n";
|
||||
|
||||
const string keyboard_shorcuts = "Keyboard Shorcuts: \n\n"
|
||||
"Press 'i' to increase brush size.\n"
|
||||
"Press 'd' to decrease brush size.\n"
|
||||
"Press 'r' to reset mask.\n"
|
||||
"Press ' ' (space bar) after selecting area to be inpainted.\n"
|
||||
"Press ESC to terminate the program.\n\n";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | show help message}"
|
||||
"{ @alias | lama | 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 | rubberwhale1.png | image file path}";
|
||||
|
||||
const string backend_keys = format(
|
||||
"{ backend | default | Choose one of computation backends: "
|
||||
"default: automatically (by default), "
|
||||
"openvino: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"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 preprocess) }");
|
||||
|
||||
string keys = param_keys + backend_keys + target_keys;
|
||||
bool drawing = false;
|
||||
Mat maskGray;
|
||||
int brush_size = 15;
|
||||
|
||||
|
||||
static void drawMask(int event, int x, int y, int, void*) {
|
||||
if (event == EVENT_LBUTTONDOWN) {
|
||||
drawing = true;
|
||||
} else if (event == EVENT_MOUSEMOVE) {
|
||||
if (drawing) {
|
||||
circle(maskGray, Point(x, y), brush_size, Scalar(255), -1);
|
||||
}
|
||||
} else if (event == EVENT_LBUTTONUP) {
|
||||
drawing = false;
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char **argv)
|
||||
{
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
|
||||
if (!parser.has("@alias") || parser.has("help"))
|
||||
{
|
||||
cout<<about<<endl;
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
string modelName = parser.get<String>("@alias");
|
||||
string zooFile = findFile(parser.get<String>("zoo"));
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
parser.about("Use this script to run image inpainting using OpenCV.");
|
||||
|
||||
const string sha1 = parser.get<String>("sha1");
|
||||
const string modelPath = findModel(parser.get<String>("model"), sha1);
|
||||
string imgPath = parser.get<String>("input");
|
||||
const string backend = parser.get<String>("backend");
|
||||
const string target = parser.get<String>("target");
|
||||
int height = parser.get<int>("height");
|
||||
int width = parser.get<int>("width");
|
||||
float scale = parser.get<float>("scale");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
Scalar mean_v = parser.get<Scalar>("mean");
|
||||
int stdSize = 20;
|
||||
int stdWeight = 400;
|
||||
int stdImgSize = 512;
|
||||
int imgWidth = -1; // Initialization
|
||||
int fontSize = 60;
|
||||
int fontWeight = 500;
|
||||
|
||||
cout<<"Model loading..."<<endl;
|
||||
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if (backend != "default" || target != "cpu"){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
|
||||
Net net = readNetFromONNX(modelPath, engine);
|
||||
net.setPreferableBackend(getBackendID(backend));
|
||||
net.setPreferableTarget(getTargetID(target));
|
||||
|
||||
FontFace fontFace("sans");
|
||||
|
||||
Mat input_image = imread(findFile(imgPath));
|
||||
if (input_image.empty()) {
|
||||
cerr << "Error: Input image could not be loaded." << endl;
|
||||
return -1;
|
||||
}
|
||||
double aspectRatio = static_cast<double>(input_image.rows) / static_cast<double>(input_image.cols);
|
||||
int h = static_cast<int>(width * aspectRatio);
|
||||
resize(input_image, input_image, Size(width, h));
|
||||
Mat image = input_image.clone();
|
||||
|
||||
imgWidth = min(input_image.rows, input_image.cols);
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize);
|
||||
fontWeight = min(fontWeight, (stdWeight*imgWidth)/stdImgSize);
|
||||
|
||||
cout<<keyboard_shorcuts<<endl;
|
||||
const string label = "Press 'i' to increase, 'd' to decrease brush size. And 'r' to reset mask. ";
|
||||
double alpha = 0.5;
|
||||
Rect r = getTextSize(Size(), label, Point(), fontFace, fontSize, fontWeight);
|
||||
r.height += 2 * fontSize; // padding
|
||||
r.width += 10; // padding
|
||||
// Setting up window
|
||||
namedWindow("Draw Mask");
|
||||
setMouseCallback("Draw Mask", drawMask);
|
||||
Mat tempImage = input_image.clone();
|
||||
Mat overlay = input_image.clone();
|
||||
rectangle(overlay, r, Scalar::all(255), FILLED);
|
||||
addWeighted(overlay, alpha, tempImage, 1 - alpha, 0, tempImage);
|
||||
putText(tempImage, "Draw the mask on the image. Press space bar when done", Point(10, fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
putText(tempImage, label, Point(10, 2*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
Mat displayImage = tempImage.clone();
|
||||
|
||||
for (;;) {
|
||||
maskGray = Mat::zeros(input_image.size(), CV_8U);
|
||||
displayImage = tempImage.clone();
|
||||
for(;;) {
|
||||
displayImage.setTo(Scalar(255, 255, 255), maskGray > 0); // Highlight mask area
|
||||
imshow("Draw Mask", displayImage);
|
||||
int key = waitKey(30) & 255;
|
||||
if (key == 'i') {
|
||||
brush_size += 1;
|
||||
cout << "Brush size increased to " << brush_size << endl;
|
||||
} else if (key == 'd') {
|
||||
brush_size = max(1, brush_size - 1);
|
||||
cout << "Brush size decreased to " << brush_size << endl;
|
||||
} else if (key == 'r') {
|
||||
maskGray = Mat::zeros(image.size(), CV_8U);
|
||||
displayImage = tempImage.clone();
|
||||
cout << "Mask cleared." << endl;
|
||||
} else if (key == ' ') {
|
||||
break;
|
||||
} else if (key == 27){
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
cout<<"Processing image..."<<endl;
|
||||
// Inference block
|
||||
Mat image_blob = blobFromImage(image, scale, Size(width, height), mean_v, swapRB, false);
|
||||
|
||||
Mat mask_blob;
|
||||
mask_blob = blobFromImage(maskGray, 1.0, Size(width, height), Scalar(0), false, false);
|
||||
mask_blob = (mask_blob > 0);
|
||||
mask_blob.convertTo(mask_blob, CV_32F);
|
||||
mask_blob = mask_blob/255.0;
|
||||
|
||||
net.setInput(image_blob, "image");
|
||||
net.setInput(mask_blob, "mask");
|
||||
|
||||
Mat output = net.forward();
|
||||
// Post Processing
|
||||
Mat output_transposed(3, &output.size[1], CV_32F, output.ptr<float>());
|
||||
|
||||
vector<Mat> channels;
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
channels.push_back(Mat(output_transposed.size[1], output_transposed.size[2], CV_32F,
|
||||
output_transposed.ptr<float>(i)));
|
||||
}
|
||||
Mat output_image;
|
||||
merge(channels, output_image);
|
||||
output_image.convertTo(output_image, CV_8U);
|
||||
|
||||
resize(output_image, output_image, Size(width, h));
|
||||
image = output_image;
|
||||
|
||||
imshow("Inpainted Output", output_image);
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
199
samples/dnn/inpainting.py
Normal file
199
samples/dnn/inpainting.py
Normal file
@ -0,0 +1,199 @@
|
||||
#!/usr/bin/env python
|
||||
'''
|
||||
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.
|
||||
|
||||
This sample inpaints the masked area in the given image.
|
||||
|
||||
Copyright (C) 2025, Bigvision LLC.
|
||||
|
||||
How to use:
|
||||
Sample command to run:
|
||||
`python inpainting.py`
|
||||
The system will ask you to draw the mask to be inpainted
|
||||
|
||||
You can download lama inpainting model using
|
||||
`python download_models.py lama`
|
||||
|
||||
References:
|
||||
Github: https://github.com/advimman/lama
|
||||
ONNX model: https://huggingface.co/Carve/LaMa-ONNX/blob/main/lama_fp32.onnx
|
||||
|
||||
ONNX model was further quantized using block quantization from [opencv_zoo](https://github.com/opencv/opencv_zoo)
|
||||
|
||||
Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to point to the directory where models are downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.
|
||||
'''
|
||||
import argparse
|
||||
import os.path
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
from common import *
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Use this script for image inpainting using OpenCV.
|
||||
|
||||
Firstly, download required models i.e. lama using `download_models.py` (if not already done). Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to specify where models should be downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.
|
||||
|
||||
To run:
|
||||
Example: python inpainting.py [--input=<image_name>]
|
||||
|
||||
Inpainting model path can also be specified using --model argument.
|
||||
'''
|
||||
)
|
||||
|
||||
def keyboard_shorcuts():
|
||||
print('''
|
||||
Keyboard Shorcuts:
|
||||
Press 'i' to increase brush size.
|
||||
Press 'd' to decrease brush size.
|
||||
Press 'r' to reset mask.
|
||||
Press ' ' (space bar) after selecting area to be inpainted.
|
||||
Press ESC to terminate the program.
|
||||
'''
|
||||
)
|
||||
|
||||
def get_args_parser():
|
||||
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', '-i', default="rubberwhale1.png", help='Path to image file.', required=False)
|
||||
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 (https://software.intel.com/openvino-toolkit), "
|
||||
"opencv: OpenCV implementation, "
|
||||
"vkcom: VKCOM, "
|
||||
"cuda: CUDA, "
|
||||
"webnn: WebNN")
|
||||
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 preprocess)")
|
||||
args, _ = parser.parse_known_args()
|
||||
add_preproc_args(args.zoo, parser, 'inpainting', prefix="", alias="lama")
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Image inpainting using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
drawing = False
|
||||
mask_gray = None
|
||||
brush_size = 15
|
||||
|
||||
def draw_mask(event, x, y, flags, param):
|
||||
global drawing, mask_gray, brush_size
|
||||
if event == cv.EVENT_LBUTTONDOWN:
|
||||
drawing = True
|
||||
elif event == cv.EVENT_MOUSEMOVE:
|
||||
if drawing:
|
||||
cv.circle(mask_gray, (x, y), brush_size, (255), thickness=-1)
|
||||
elif event == cv.EVENT_LBUTTONUP:
|
||||
drawing = False
|
||||
|
||||
def main():
|
||||
global mask_gray, brush_size
|
||||
|
||||
print("Model loading...")
|
||||
|
||||
if hasattr(args, 'help'):
|
||||
help()
|
||||
exit(1)
|
||||
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
|
||||
engine = cv.dnn.ENGINE_AUTO
|
||||
|
||||
if args.backend != "default" or args.target != "cpu":
|
||||
engine = cv.dnn.ENGINE_CLASSIC
|
||||
|
||||
net = cv.dnn.readNetFromONNX(args.model, engine)
|
||||
net.setPreferableBackend(get_backend_id(args.backend))
|
||||
net.setPreferableTarget(get_target_id(args.target))
|
||||
|
||||
input_image = cv.imread(findFile(args.input))
|
||||
aspect_ratio = input_image.shape[0]/input_image.shape[1]
|
||||
height = int(args.width*aspect_ratio)
|
||||
|
||||
input_image = cv.resize(input_image, (args.width, height))
|
||||
image = input_image.copy()
|
||||
keyboard_shorcuts()
|
||||
|
||||
stdSize = 0.7
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
imgWidth = min(input_image.shape[:2])
|
||||
fontSize = min(1.5, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(1,(stdWeight*imgWidth)//stdImgSize)
|
||||
|
||||
label = "Press 'i' to increase, 'd' to decrease brush size. And 'r' to reset mask. "
|
||||
labelSize, _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
alpha = 0.5
|
||||
# Setting up the window
|
||||
cv.namedWindow("Draw Mask")
|
||||
cv.setMouseCallback("Draw Mask", draw_mask)
|
||||
temp_image = input_image.copy()
|
||||
overlay = input_image.copy()
|
||||
cv.rectangle(overlay, (0, 0), (labelSize[0]+10, labelSize[1]+int(30*fontSize)), (255, 255, 255), cv.FILLED)
|
||||
cv.addWeighted(overlay, alpha, temp_image, 1 - alpha, 0, temp_image)
|
||||
cv.putText(temp_image, "Draw the mask on the image. Press space bar when done.", (10, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.putText(temp_image, label, (10, int(50*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
display_image = temp_image.copy()
|
||||
|
||||
while True:
|
||||
mask_gray = np.zeros((input_image.shape[0], input_image.shape[1]), dtype=np.uint8)
|
||||
display_image = temp_image.copy()
|
||||
while True:
|
||||
display_image[mask_gray > 0] = [255, 255, 255]
|
||||
cv.imshow("Draw Mask", display_image)
|
||||
key = cv.waitKey(30) & 0xFF
|
||||
if key == ord('i'): # Increase brush size
|
||||
brush_size += 1
|
||||
print(f"Brush size increased to {brush_size}")
|
||||
elif key == ord('d'): # Decrease brush size
|
||||
brush_size = max(1, brush_size - 1)
|
||||
print(f"Brush size decreased to {brush_size}")
|
||||
elif key == ord('r'): # clear the mask
|
||||
mask_gray = np.zeros((input_image.shape[0], input_image.shape[1]), dtype=np.uint8)
|
||||
display_image = temp_image.copy()
|
||||
print(f"Mask cleared")
|
||||
elif key == ord(' '): # Press space bar to finish drawing
|
||||
break
|
||||
elif key == 27:
|
||||
exit()
|
||||
|
||||
print("Processing image...")
|
||||
# Inference block
|
||||
image_blob = cv.dnn.blobFromImage(image, args.scale, (args.width, args.height), args.mean, args.rgb, False)
|
||||
mask_blob = cv.dnn.blobFromImage(mask_gray, scalefactor=1.0, size=(args.width, args.height), mean=(0,), swapRB=False, crop=False)
|
||||
mask_blob = (mask_blob > 0).astype(np.float32)
|
||||
|
||||
net.setInput(image_blob, "image")
|
||||
net.setInput(mask_blob, "mask")
|
||||
|
||||
output = net.forward()
|
||||
|
||||
# Postprocessing
|
||||
output_image = output[0]
|
||||
output_image = np.transpose(output_image, (1, 2, 0))
|
||||
output_image = (output_image).astype(np.uint8)
|
||||
output_image = cv.resize(output_image, (args.width, height))
|
||||
image = output_image
|
||||
|
||||
cv.imshow("Inpainted Output", output_image)
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args_parser()
|
||||
main()
|
||||
@ -440,4 +440,20 @@ dasiamrpn:
|
||||
dasiamrpn_kernel_cls_url: "https://github.com/opencv/opencv_zoo/raw/fef72f8fa7c52eaf116d3df358d24e6e959ada0e/models/object_tracking_dasiamrpn/object_tracking_dasiamrpn_kernel_cls1_2021nov.onnx?download="
|
||||
dasiamrpn_kernel_cls_sha1: "e9ccd270ce8059bdf7ed0d1845c03ef4a951ee0f"
|
||||
dasiamrpn_kernel_cls_model: "object_tracking_dasiamrpn_kernel_cls1_2021nov.onnx"
|
||||
sample: "object_tracker"
|
||||
sample: "object_tracker"
|
||||
|
||||
################################################################################
|
||||
# Inpainting models.
|
||||
################################################################################
|
||||
|
||||
lama:
|
||||
load_info:
|
||||
url: "https://github.com/gursimarsingh/opencv_zoo/raw/0417e12d24bba41613ae0380bd698cca73a4fb17/models/inpainting_lama/inpainting_lama_2025jan.onnx?download="
|
||||
sha1: "7c6cdb9362bf73de2a80cfcaf17e121e3302f24c"
|
||||
model: "inpainting_lama_2025jan.onnx"
|
||||
mean: [0, 0, 0]
|
||||
scale: 0.00392
|
||||
width: 512
|
||||
height: 512
|
||||
rgb: false
|
||||
sample: "inpainting"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user