vendor: OpenCV 5.0.0 snapshot at 40738fb16ceddb5fb3fea747585f7ce6abb0605b
This commit is contained in:
Gitea Mirror Bot
@@ -0,0 +1,3 @@
|
||||
*.caffemodel
|
||||
*.pb
|
||||
*.weights
|
||||
@@ -0,0 +1,24 @@
|
||||
ocv_install_example_src(dnn *.cpp *.hpp CMakeLists.txt)
|
||||
|
||||
set(OPENCV_DNN_SAMPLES_REQUIRED_DEPS
|
||||
opencv_core
|
||||
opencv_imgproc
|
||||
opencv_dnn
|
||||
opencv_objdetect
|
||||
opencv_video
|
||||
opencv_imgcodecs
|
||||
opencv_videoio
|
||||
opencv_highgui)
|
||||
ocv_check_dependencies(${OPENCV_DNN_SAMPLES_REQUIRED_DEPS})
|
||||
|
||||
if(NOT BUILD_EXAMPLES OR NOT OCV_DEPENDENCIES_FOUND)
|
||||
return()
|
||||
endif()
|
||||
|
||||
project(dnn_samples)
|
||||
ocv_include_modules_recurse(${OPENCV_DNN_SAMPLES_REQUIRED_DEPS})
|
||||
file(GLOB_RECURSE dnn_samples RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} *.cpp)
|
||||
foreach(sample_filename ${dnn_samples})
|
||||
ocv_define_sample(tgt ${sample_filename} dnn)
|
||||
ocv_target_link_libraries(${tgt} PRIVATE ${OPENCV_LINKER_LIBS} ${OPENCV_DNN_SAMPLES_REQUIRED_DEPS})
|
||||
endforeach()
|
||||
@@ -0,0 +1,84 @@
|
||||
# OpenCV deep learning module samples
|
||||
|
||||
## Model Zoo
|
||||
|
||||
Check [a wiki](https://github.com/opencv/opencv/wiki/Deep-Learning-in-OpenCV) for a list of tested models.
|
||||
|
||||
If OpenCV is built with [Intel's Inference Engine support](https://github.com/opencv/opencv/wiki/Intel%27s-Deep-Learning-Inference-Engine-backend) you can use [Intel's pre-trained](https://github.com/opencv/open_model_zoo) models.
|
||||
|
||||
There are different preprocessing parameters such mean subtraction or scale factors for different models.
|
||||
You may check the most popular models and their parameters at [models.yml](https://github.com/opencv/opencv/blob/5.x/samples/dnn/models.yml) configuration file. It might be also used for aliasing samples parameters. In example,
|
||||
|
||||
```bash
|
||||
python object_detection.py opencv_fd --model /path/to/model.onnx
|
||||
```
|
||||
|
||||
Check `-h` option to know which values are used by default:
|
||||
|
||||
```bash
|
||||
python object_detection.py opencv_fd -h
|
||||
```
|
||||
|
||||
### Sample models
|
||||
|
||||
You can download sample models using ```download_models.py```. For example, the following command will download network weights for OpenCV Face Detector model and store them in FaceDetector folder:
|
||||
|
||||
```bash
|
||||
python download_models.py --save_dir FaceDetector opencv_fd
|
||||
```
|
||||
|
||||
You can use default configuration files adopted for OpenCV from [here](https://github.com/opencv/opencv_extra/tree/5.x/testdata/dnn).
|
||||
|
||||
You also can use the script to download necessary files from your code. Assume you have the following code inside ```your_script.py```:
|
||||
|
||||
```python
|
||||
from download_models import downloadFile
|
||||
|
||||
filepath1 = downloadFile("https://huggingface.co/onnxmodelzoo/ssd_mobilenet_v1_12/resolve/main/ssd_mobilenet_v1_12.onnx", None, filename="ssd_mobilenet_v1_12.onnx", save_dir="save_dir_1")
|
||||
filepath2 = downloadFile("https://huggingface.co/onnxmodelzoo/ssd_mobilenet_v1_12/resolve/main/ssd_mobilenet_v1_12.onnx", "83536889adce1eda154175f8e3b156dd20443631", filename="ssd_mobilenet_v1_12.onnx")
|
||||
print(filepath1)
|
||||
print(filepath2)
|
||||
# Your code
|
||||
```
|
||||
|
||||
By running the following commands, you will get **ssd_mobilenet_v1_12.onnx** file:
|
||||
```bash
|
||||
export OPENCV_DOWNLOAD_DATA_PATH=download_folder
|
||||
python your_script.py
|
||||
```
|
||||
|
||||
**Note** that you can provide a directory using **save_dir** parameter or via **OPENCV_SAVE_DIR** environment variable.
|
||||
|
||||
#### Face detection
|
||||
[An origin model](https://github.com/opencv/opencv/tree/5.x/samples/dnn/face_detector)
|
||||
with single precision floating point weights has been quantized using [TensorFlow framework](https://www.tensorflow.org/).
|
||||
To achieve the best accuracy run the model on BGR images resized to `300x300` applying mean subtraction
|
||||
of values `(104, 177, 123)` for each blue, green and red channels correspondingly.
|
||||
|
||||
The following are accuracy metrics obtained using [COCO object detection evaluation
|
||||
tool](http://cocodataset.org/#detections-eval) on [FDDB dataset](http://vis-www.cs.umass.edu/fddb/)
|
||||
(see [script](https://github.com/opencv/opencv/blob/5.x/modules/dnn/misc/face_detector_accuracy.py))
|
||||
applying resize to `300x300` and keeping an origin images' sizes.
|
||||
```
|
||||
AP - Average Precision | FP32/FP16 | UINT8 | FP32/FP16 | UINT8 |
|
||||
AR - Average Recall | 300x300 | 300x300 | any size | any size |
|
||||
--------------------------------------------------|-----------|----------------|-----------|----------------|
|
||||
AP @[ IoU=0.50:0.95 | area= all | maxDets=100 ] | 0.408 | 0.408 | 0.378 | 0.328 (-0.050) |
|
||||
AP @[ IoU=0.50 | area= all | maxDets=100 ] | 0.849 | 0.849 | 0.797 | 0.790 (-0.007) |
|
||||
AP @[ IoU=0.75 | area= all | maxDets=100 ] | 0.251 | 0.251 | 0.208 | 0.140 (-0.068) |
|
||||
AP @[ IoU=0.50:0.95 | area= small | maxDets=100 ] | 0.050 | 0.051 (+0.001) | 0.107 | 0.070 (-0.037) |
|
||||
AP @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] | 0.381 | 0.379 (-0.002) | 0.380 | 0.368 (-0.012) |
|
||||
AP @[ IoU=0.50:0.95 | area= large | maxDets=100 ] | 0.455 | 0.455 | 0.412 | 0.337 (-0.075) |
|
||||
AR @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] | 0.299 | 0.299 | 0.279 | 0.246 (-0.033) |
|
||||
AR @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] | 0.482 | 0.482 | 0.476 | 0.436 (-0.040) |
|
||||
AR @[ IoU=0.50:0.95 | area= all | maxDets=100 ] | 0.496 | 0.496 | 0.491 | 0.451 (-0.040) |
|
||||
AR @[ IoU=0.50:0.95 | area= small | maxDets=100 ] | 0.189 | 0.193 (+0.004) | 0.284 | 0.232 (-0.052) |
|
||||
AR @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] | 0.481 | 0.480 (-0.001) | 0.470 | 0.458 (-0.012) |
|
||||
AR @[ IoU=0.50:0.95 | area= large | maxDets=100 ] | 0.528 | 0.528 | 0.520 | 0.462 (-0.058) |
|
||||
```
|
||||
|
||||
## References
|
||||
* [Models downloading script](https://github.com/opencv/opencv/blob/5.x/samples/dnn/download_models.py)
|
||||
* [Configuration files adopted for OpenCV](https://github.com/opencv/opencv_extra/tree/5.x/testdata/dnn)
|
||||
* [How to import models from TensorFlow Object Detection API](https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API)
|
||||
* [Names of classes from different datasets](https://github.com/opencv/opencv/tree/5.x/samples/data/dnn)
|
||||
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
from common import findFile
|
||||
|
||||
parser = argparse.ArgumentParser(description='Use this script to run action recognition using 3D ResNet34',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input', '-i', help='Path to input video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--model', required=True, help='Path to model.')
|
||||
parser.add_argument('--classes', default=findFile('action_recongnition_kinetics.txt'), help='Path to classes list.')
|
||||
|
||||
# To get net download original repository https://github.com/kenshohara/video-classification-3d-cnn-pytorch
|
||||
# For correct ONNX export modify file: video-classification-3d-cnn-pytorch/models/resnet.py
|
||||
# change
|
||||
# - def downsample_basic_block(x, planes, stride):
|
||||
# - out = F.avg_pool3d(x, kernel_size=1, stride=stride)
|
||||
# - zero_pads = torch.Tensor(out.size(0), planes - out.size(1),
|
||||
# - out.size(2), out.size(3),
|
||||
# - out.size(4)).zero_()
|
||||
# - if isinstance(out.data, torch.cuda.FloatTensor):
|
||||
# - zero_pads = zero_pads.cuda()
|
||||
# -
|
||||
# - out = Variable(torch.cat([out.data, zero_pads], dim=1))
|
||||
# - return out
|
||||
|
||||
# To
|
||||
# + def downsample_basic_block(x, planes, stride):
|
||||
# + out = F.avg_pool3d(x, kernel_size=1, stride=stride)
|
||||
# + out = F.pad(out, (0, 0, 0, 0, 0, 0, 0, int(planes - out.size(1)), 0, 0), "constant", 0)
|
||||
# + return out
|
||||
|
||||
# To ONNX export use torch.onnx.export(model, inputs, model_name)
|
||||
|
||||
def get_class_names(path):
|
||||
class_names = []
|
||||
with open(path) as f:
|
||||
for row in f:
|
||||
class_names.append(row[:-1])
|
||||
return class_names
|
||||
|
||||
def classify_video(video_path, net_path):
|
||||
SAMPLE_DURATION = 16
|
||||
SAMPLE_SIZE = 112
|
||||
mean = (114.7748, 107.7354, 99.4750)
|
||||
class_names = get_class_names(args.classes)
|
||||
|
||||
net = cv.dnn.readNet(net_path)
|
||||
net.setPreferableBackend(cv.dnn.DNN_BACKEND_INFERENCE_ENGINE)
|
||||
net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)
|
||||
|
||||
winName = 'Deep learning image classification in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_AUTOSIZE)
|
||||
cap = cv.VideoCapture(video_path)
|
||||
while cv.waitKey(1) < 0:
|
||||
frames = []
|
||||
for _ in range(SAMPLE_DURATION):
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
exit(0)
|
||||
frames.append(frame)
|
||||
|
||||
inputs = cv.dnn.blobFromImages(frames, 1, (SAMPLE_SIZE, SAMPLE_SIZE), mean, True, crop=True)
|
||||
inputs = np.transpose(inputs, (1, 0, 2, 3))
|
||||
inputs = np.expand_dims(inputs, axis=0)
|
||||
net.setInput(inputs)
|
||||
outputs = net.forward()
|
||||
class_pred = np.argmax(outputs)
|
||||
label = class_names[class_pred]
|
||||
|
||||
for frame in frames:
|
||||
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||
cv.rectangle(frame, (0, 10 - labelSize[1]),
|
||||
(labelSize[0], 10 + baseLine), (255, 255, 255), cv.FILLED)
|
||||
cv.putText(frame, label, (0, 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
|
||||
cv.imshow(winName, frame)
|
||||
if cv.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
if __name__ == "__main__":
|
||||
args, _ = parser.parse_known_args()
|
||||
classify_video(args.input if args.input else 0, args.model)
|
||||
@@ -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 <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
#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<String>("@alias");
|
||||
string zooFile = parser.get<String>("zoo");
|
||||
|
||||
zooFile = findFile(zooFile);
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
|
||||
int input_width = parser.get<int>("width");
|
||||
int input_height = parser.get<int>("height");
|
||||
float scale_factor = parser.get<float>("scale");
|
||||
Scalar mean_values = parser.get<Scalar>("mean");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
String backend = parser.get<String>("backend");
|
||||
String target = parser.get<String>("target");
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
|
||||
string model = findModel(parser.get<String>("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<String>("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;
|
||||
}
|
||||
@@ -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)
|
||||
t0 = cv.getTickCount()
|
||||
out = model.forward()
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
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, t
|
||||
|
||||
|
||||
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, t = apply_modnet(args, net, image)
|
||||
label = "Inference time: %.2f ms" % (t * 1000.0)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,265 @@
|
||||
// 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.
|
||||
|
||||
/*
|
||||
Auto white balance using FC4: https://github.com/yuanming-hu/fc4
|
||||
|
||||
Color constancy is a method to make colors of objects render correctly on a photo.
|
||||
White balance aims to make white objects appear white on an image and not a shade of any
|
||||
other color, independent of the actual light setting. White balance correction creates
|
||||
a neutral looking coloring of the objects, and generally makes colors look more similar
|
||||
to their 'true' colors under different light conditions.
|
||||
|
||||
Given an RGB image, the FC4 model predicts scene illuminant (R,G,B). We then apply
|
||||
the illuminant to the image, applying the correction in the linear RGB space.
|
||||
The transformation between linear and sRGB spaces is done as described in the sRGB standard,
|
||||
which is a nonlinear Gamma correction with exponent 2.4 and extra handling of very small values.
|
||||
This sample is written for 8bit images. The FC4 model accepts RGB images with applied Gamma scaling.
|
||||
|
||||
The training of the FC4 model was done on the Gehler-Shi dataset. The dataset includes
|
||||
568 images and ground truth corrections, as well as ground truth illuminants. The linear
|
||||
RGB images from the dataset were used with Gamma correction of 2.2 applied.
|
||||
|
||||
The model is a pretrained fold 0 of a training pipeline on the Gehler-Shi dataset, from the PyTorch
|
||||
implementation of the FC4 algorithm by Mateo Rizzo. The model was converted from a .pth file to onnx
|
||||
using torch.onnx.export. The model can be downloaded in the following link:
|
||||
https://raw.githubusercontent.com/MykhailoTrushch/opencv/d6ab21353a87e4c527e38e464384c7ee78e96e22/samples/dnn/models/fc4_fold_0.onnx
|
||||
|
||||
Copyright (c) 2017 Yuanming Hu, Baoyuan Wang, Stephen Lin
|
||||
Copyright (c) 2021 Matteo Rizzo
|
||||
|
||||
Licensed under the MIT license.
|
||||
|
||||
References:
|
||||
|
||||
Yuanming Hu, Baoyuan Wang, and Stephen Lin. “FC⁴: Fully Convolutional Color
|
||||
Constancy with Confidence-Weighted Pooling.” CVPR, 2017, pp. 4085–4094.
|
||||
|
||||
Implementations of FC4:
|
||||
https://github.com/yuanming-hu/fc4/
|
||||
https://github.com/matteo-rizzo/fc4-pytorch
|
||||
|
||||
Lilong Shi and Brian Funt, "Re-processed Version of the Gehler Color Constancy Dataset of 568 Images,"
|
||||
accessed from http://www.cs.sfu.ca/~colour/data/
|
||||
|
||||
“IEC 61966-2-1:1999 – Multimedia Systems and Equipment – Colour Measurement and Management – Part 2-1: Colour Management – Default RGB Colour Space – sRGB.” IEC Standard, 1999.
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace cv::dnn;
|
||||
using namespace std;
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message }"
|
||||
"{ @alias | fc4 | Model alias from models.yml "
|
||||
"(optional) }"
|
||||
"{ zoo | ../dnn/models.yml | Path to models.yml file "
|
||||
"(optional) }"
|
||||
"{ input i | castle.png | Path to input image }";
|
||||
;
|
||||
|
||||
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) }");
|
||||
|
||||
// Normalization constant for 8bit values
|
||||
const float NORMALIZE_FACTOR = 1.0f / 255.0f;
|
||||
|
||||
// sRGB to linear conversion constants (or vice versa):
|
||||
// SRGB_THRESHOLD / LINEAR_THRESHOLD: breakpoints between linear and gamma regions
|
||||
// SRGB_SLOPE: slope of the linear segment near black
|
||||
// SRGB_ALPHA: offset to ensure continuity at the threshold
|
||||
// SRGB_EXP: gamma exponent
|
||||
const float SRGB_THRESHOLD = 0.04045f;
|
||||
const float SRGB_ALPHA = 0.055f;
|
||||
const float SRGB_SLOPE = 12.92f;
|
||||
const float SRGB_EXP = 2.4f;
|
||||
const float LINEAR_THRESHOLD = 0.0031308f;
|
||||
const float EPS = 1e-10f;
|
||||
|
||||
static Mat srgbToLinear(const Mat &srgb32f) {
|
||||
CV_Assert(srgb32f.type() == CV_32FC3);
|
||||
const float a = SRGB_ALPHA;
|
||||
|
||||
Mat y = srgb32f;
|
||||
Mat mask_low;
|
||||
compare(y, SRGB_THRESHOLD, mask_low, CMP_LE);
|
||||
|
||||
Mat low = y / SRGB_SLOPE;
|
||||
|
||||
Mat t = (y + a) / (1.0f + a);
|
||||
Mat high;
|
||||
pow(t, SRGB_EXP, high);
|
||||
|
||||
Mat lin(y.size(), y.type(), Scalar(0, 0, 0));
|
||||
low.copyTo(lin, mask_low);
|
||||
Mat mask_high;
|
||||
bitwise_not(mask_low, mask_high);
|
||||
high.copyTo(lin, mask_high);
|
||||
|
||||
return lin;
|
||||
}
|
||||
|
||||
static Mat linearToSrgb(const Mat &lin32f) {
|
||||
CV_Assert(lin32f.type() == CV_32FC3);
|
||||
const float a = SRGB_ALPHA;
|
||||
|
||||
Mat x = lin32f;
|
||||
Mat mask_low;
|
||||
compare(x, LINEAR_THRESHOLD, mask_low, CMP_LE);
|
||||
Mat low = x * SRGB_SLOPE;
|
||||
Mat powPart;
|
||||
pow(x, 1.0 / SRGB_EXP, powPart);
|
||||
Mat high = (1.0f + a) * powPart - a;
|
||||
|
||||
Mat srgb(x.size(), x.type(), Scalar(0, 0, 0));
|
||||
low.copyTo(srgb, mask_low);
|
||||
Mat mask_high;
|
||||
bitwise_not(mask_low, mask_high);
|
||||
high.copyTo(srgb, mask_high);
|
||||
|
||||
return srgb;
|
||||
}
|
||||
|
||||
static Mat correct(const Mat &bgr8u, const Vec3f &illumRGB_linear) {
|
||||
Mat f32;
|
||||
bgr8u.convertTo(f32, CV_32F, NORMALIZE_FACTOR);
|
||||
|
||||
Mat lin = srgbToLinear(f32);
|
||||
|
||||
const float eR = std::max(illumRGB_linear[0], EPS);
|
||||
const float eG = std::max(illumRGB_linear[1], EPS);
|
||||
const float eB = std::max(illumRGB_linear[2], EPS);
|
||||
|
||||
float s3 = std::sqrt(3.0f);
|
||||
Scalar corr(eB * s3 + EPS, eG * s3 + EPS, eR * s3 + EPS);
|
||||
|
||||
Mat corrected;
|
||||
divide(lin, corr, corrected);
|
||||
|
||||
std::vector<Mat> ch;
|
||||
split(corrected, ch);
|
||||
double m0, m1, m2;
|
||||
minMaxLoc(ch[0], nullptr, &m0);
|
||||
minMaxLoc(ch[1], nullptr, &m1);
|
||||
minMaxLoc(ch[2], nullptr, &m2);
|
||||
float maxVal = static_cast<float>(std::max({m0, m1, m2})) + EPS;
|
||||
corrected /= maxVal;
|
||||
min(corrected, 1.0, corrected);
|
||||
max(corrected, 0.0, corrected);
|
||||
|
||||
Mat srgb = linearToSrgb(corrected);
|
||||
|
||||
Mat out;
|
||||
srgb.convertTo(out, CV_8U, 255.0);
|
||||
return out;
|
||||
}
|
||||
|
||||
static void annotate(Mat &img, const string &title) {
|
||||
double fs = std::max(0.5, std::min(img.cols, img.rows) / 800.0);
|
||||
int th = std::max(1, (int)std::round(fs * 2));
|
||||
putText(img, title, Point(10, 30), FONT_HERSHEY_SIMPLEX, fs,
|
||||
Scalar(0, 255, 0), th);
|
||||
}
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
const string about = "FC4 Color Constancy (ONNX) sample.\n"
|
||||
"Predicts scene illuminant and corrects the white "
|
||||
"balance of the image.\n";
|
||||
|
||||
string keys = param_keys + backend_keys + target_keys;
|
||||
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
if (parser.has("help")) {
|
||||
cout << about << endl;
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
string modelName = parser.get<String>("@alias");
|
||||
string zooFile = samples::findFile(parser.get<String>("zoo"));
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
|
||||
float scale = parser.get<float>("scale");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
String backend = parser.get<String>("backend");
|
||||
String target = parser.get<String>("target");
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
string model = findModel(parser.get<String>("model"), sha1);
|
||||
string inputPath = findFile(parser.get<String>("input"));
|
||||
|
||||
if (model.empty()) {
|
||||
cerr << "Model file not found\n";
|
||||
return -1;
|
||||
}
|
||||
|
||||
Net net;
|
||||
try {
|
||||
net = readNetFromONNX(model);
|
||||
net.setPreferableBackend(getBackendID(backend));
|
||||
net.setPreferableTarget(getTargetID(target));
|
||||
} catch (const Exception &e) {
|
||||
cerr << "Error loading model: " << e.what() << endl;
|
||||
return -1;
|
||||
}
|
||||
Mat img = imread(inputPath, IMREAD_COLOR);
|
||||
if (img.empty()) {
|
||||
cerr << "Cannot load image: " << inputPath << endl;
|
||||
return -1;
|
||||
}
|
||||
Mat blob;
|
||||
blob = blobFromImage(img, scale, img.size(), mean, swapRB, /*crop=*/false,
|
||||
/*type=*/CV_32F);
|
||||
net.setInput(blob);
|
||||
|
||||
Mat out;
|
||||
try {
|
||||
out = net.forward();
|
||||
} catch (const Exception &e) {
|
||||
cerr << "Forward error: " << e.what() << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
const float *p = out.ptr<float>(0);
|
||||
CV_Assert(out.total() == 3);
|
||||
Vec3f illum = Vec3f(p[0], p[1], p[2]);
|
||||
|
||||
Mat corrected = correct(img, illum);
|
||||
|
||||
Mat origVis = img.clone();
|
||||
Mat corrVis = corrected.clone();
|
||||
annotate(origVis, "Original");
|
||||
annotate(corrVis, "FC4-corrected");
|
||||
Mat stacked;
|
||||
hconcat(origVis, corrVis, stacked);
|
||||
imshow("Original and Corrected Images", stacked);
|
||||
waitKey(0);
|
||||
destroyAllWindows();
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
#!/usr/bin/env python3
|
||||
# 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.
|
||||
|
||||
'''
|
||||
Auto white balance using FC4: https://github.com/yuanming-hu/fc4
|
||||
|
||||
Color constancy is a method to make colors of objects render correctly on a photo.
|
||||
White balance aims to make white objects appear white on an image and not a shade of any
|
||||
other color, independent of the actual light setting. White balance correction creates
|
||||
a neutral looking coloring of the objects, and generally makes colors look more similar
|
||||
to their 'true' colors under different light conditions.
|
||||
|
||||
Given an RGB image, the FC4 model predicts scene illuminant (R,G,B). We then apply
|
||||
the illuminant to the image, applying the correction in the linear RGB space.
|
||||
The transformation between linear and sRGB spaces is done as described in the sRGB standard,
|
||||
which is a nonlinear Gamma correction with exponent 2.4 and extra handling of very small values.
|
||||
This sample is written for 8bit images. The FC4 model accepts RGB images with applied Gamma scaling.
|
||||
|
||||
The training of the FC4 model was done on the Gehler-Shi dataset. The dataset includes
|
||||
568 images and ground truth corrections, as well as ground truth illuminants. The linear
|
||||
RGB images from the dataset were used with Gamma correction of 2.2 applied.
|
||||
|
||||
The model is a pretrained fold 0 of a training pipeline on the Gehler-Shi dataset, from the PyTorch
|
||||
implementation of the FC4 algorithm by Mateo Rizzo. The model was converted from a .pth file to onnx
|
||||
using torch.onnx.export. The model can be downloaded in the following link:
|
||||
https://raw.githubusercontent.com/MykhailoTrushch/opencv/d6ab21353a87e4c527e38e464384c7ee78e96e22/samples/dnn/models/fc4_fold_0.onnx
|
||||
|
||||
Copyright (c) 2017 Yuanming Hu, Baoyuan Wang, Stephen Lin
|
||||
Copyright (c) 2021 Matteo Rizzo
|
||||
|
||||
Licensed under the MIT license.
|
||||
|
||||
References:
|
||||
|
||||
Yuanming Hu, Baoyuan Wang, and Stephen Lin. “FC⁴: Fully Convolutional Color
|
||||
Constancy with Confidence-Weighted Pooling.” CVPR, 2017, pp. 4085–4094.
|
||||
|
||||
Implementations of FC4:
|
||||
https://github.com/yuanming-hu/fc4/
|
||||
https://github.com/matteo-rizzo/fc4-pytorch
|
||||
|
||||
Lilong Shi and Brian Funt, "Re-processed Version of the Gehler Color
|
||||
Constancy Dataset of 568 Images," accessed from http://www.cs.sfu.ca/~colour/data/
|
||||
|
||||
“IEC 61966-2-1:1999 – Multimedia Systems and Equipment – Colour Measurement and Management –
|
||||
Part 2-1: Colour Management – Default RGB Colour Space – sRGB.” IEC Standard, 1999.
|
||||
'''
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
|
||||
from common import *
|
||||
|
||||
|
||||
# Normalization constant for 8bit values
|
||||
NORMALIZE_FACTOR = 1.0 / 255.0
|
||||
|
||||
# sRGB to linear conversion constants (or vice versa):
|
||||
# SRGB_THRESHOLD / LINEAR_THRESHOLD: breakpoints between linear and gamma regions
|
||||
# SRGB_SLOPE: slope of the linear segment near black
|
||||
# SRGB_ALPHA: offset to ensure continuity at the threshold
|
||||
# SRGB_EXP: gamma exponent
|
||||
SRGB_THRESHOLD = 0.04045
|
||||
SRGB_ALPHA = 0.055
|
||||
SRGB_SLOPE = 12.92
|
||||
SRGB_EXP = 2.4
|
||||
LINEAR_THRESHOLD = 0.0031308
|
||||
EPS = 1e-10
|
||||
|
||||
def srgb_to_linear(rgb: np.ndarray) -> np.ndarray:
|
||||
low = rgb / SRGB_SLOPE
|
||||
high = np.power((rgb + SRGB_ALPHA) / (1.0 + SRGB_ALPHA), SRGB_EXP, dtype=np.float32)
|
||||
return np.where(rgb <= SRGB_THRESHOLD, low, high).astype(np.float32)
|
||||
|
||||
def linear_to_srgb(lin: np.ndarray) -> np.ndarray:
|
||||
low = lin * SRGB_SLOPE
|
||||
high = (1.0 + SRGB_ALPHA) * np.power(lin, 1.0 / SRGB_EXP, dtype=np.float32) - SRGB_ALPHA
|
||||
return np.where(lin <= LINEAR_THRESHOLD, low, high).astype(np.float32)
|
||||
|
||||
def correct(bgr8u: np.ndarray, illum_rgb_linear: np.ndarray) -> np.ndarray:
|
||||
assert bgr8u.dtype == np.uint8 and bgr8u.ndim == 3 and bgr8u.shape[2] == 3
|
||||
|
||||
bgr = bgr8u.astype(np.float32) * NORMALIZE_FACTOR
|
||||
lin = srgb_to_linear(bgr)
|
||||
e_r = max(float(illum_rgb_linear[0]), EPS)
|
||||
e_g = max(float(illum_rgb_linear[1]), EPS)
|
||||
e_b = max(float(illum_rgb_linear[2]), EPS)
|
||||
s3 = np.float32(np.sqrt(3.0))
|
||||
corr_bgr = np.array([e_b * s3 + EPS,
|
||||
e_g * s3 + EPS,
|
||||
e_r * s3 + EPS],
|
||||
dtype=np.float32)
|
||||
|
||||
corrected = lin / corr_bgr.reshape(1, 1, 3)
|
||||
|
||||
max_val = float(corrected.max()) + EPS
|
||||
corrected /= max_val
|
||||
corrected = np.clip(corrected, 0.0, 1.0)
|
||||
|
||||
srgb = linear_to_srgb(corrected)
|
||||
|
||||
out_bgr8 = (srgb * 255.0 + 0.5).astype(np.uint8)
|
||||
return out_bgr8
|
||||
|
||||
def annotate(img_bgr: np.ndarray, title: str) -> None:
|
||||
fs = max(0.5, min(img_bgr.shape[1], img_bgr.shape[0]) / 800.0)
|
||||
th = max(1, int(round(fs * 2)))
|
||||
cv.putText(img_bgr, title, (10, 30), cv.FONT_HERSHEY_SIMPLEX, fs, (0,255,0), th)
|
||||
|
||||
def get_args_parser(func_args):
|
||||
backends = ("default", "openvino", "opencv", "vkcom", "cuda", "webnn")
|
||||
targets = ("cpu", "opencl", "opencl_fp16", "ncs2_vpu", "hddl_vpu", "vulkan",
|
||||
"cuda", "cuda_fp16")
|
||||
|
||||
p = argparse.ArgumentParser(add_help=False)
|
||||
p.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.')
|
||||
p.add_argument("--input", help="Path to input image", default="castle.png")
|
||||
p.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")
|
||||
p.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, _ = p.parse_known_args()
|
||||
add_preproc_args(args.zoo, p, 'auto_white_balance', prefix="", alias="fc4")
|
||||
p = argparse.ArgumentParser(
|
||||
parents=[p],
|
||||
description="FC4 Color Constancy (ONNX): " \
|
||||
"predicts illuminant and applies white balance.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
return p.parse_args(func_args)
|
||||
|
||||
|
||||
|
||||
def main(func_args=None):
|
||||
args = get_args_parser(func_args)
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
|
||||
try:
|
||||
net = cv.dnn.readNetFromONNX(args.model)
|
||||
net.setPreferableBackend(get_backend_id(args.backend))
|
||||
net.setPreferableTarget(get_target_id(args.target))
|
||||
except cv.error as e:
|
||||
print(f"Error loading model: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
img = cv.imread(findFile(args.input), cv.IMREAD_COLOR)
|
||||
if img is None:
|
||||
print(f"Cannot load image: {args.input}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
blob = cv.dnn.blobFromImage(
|
||||
img, scalefactor=args.scale, size=(img.shape[1], img.shape[0]),
|
||||
mean=args.mean, swapRB=args.rgb, crop=False, ddepth=cv.CV_32F
|
||||
)
|
||||
net.setInput(blob)
|
||||
|
||||
try:
|
||||
out = net.forward()
|
||||
except cv.error as e:
|
||||
print(f"Forward error: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
illum = out.astype(np.float32).reshape(-1)
|
||||
if out.size != 3:
|
||||
print("Error: model output of size not equal to 3 (should output 3 illuminants in RGB order)")
|
||||
sys.exit(-1)
|
||||
|
||||
corrected = correct(img, illum)
|
||||
|
||||
orig_vis = img.copy()
|
||||
corr_vis = corrected.copy()
|
||||
annotate(orig_vis, "Original")
|
||||
annotate(corr_vis, "FC4-corrected")
|
||||
stacked = np.hstack([orig_vis, corr_vis])
|
||||
cv.imshow("Original and Corrected Images", stacked)
|
||||
cv.waitKey(0)
|
||||
cv.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,272 @@
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
#include <iostream>
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/core/utils/logger.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
using namespace dnn;
|
||||
|
||||
const string about =
|
||||
"Use this script to run a classification model on a camera stream, video, image or image list (i.e. .xml or .yaml containing image lists)\n\n"
|
||||
"Firstly, download required models 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.\n"
|
||||
"To run:\n"
|
||||
"\t ./example_dnn_classification model_name --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)\n"
|
||||
"Sample command:\n"
|
||||
"\t ./example_dnn_classification resnet --input=$OPENCV_SAMPLES_DATA_PATH/baboon.jpg\n"
|
||||
"\t ./example_dnn_classification squeezenet\n"
|
||||
"Model path can also be specified using --model argument. "
|
||||
"Use imagelist_creator to create the xml or yaml list\n";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ @alias | | 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 | | Path to input image or video file. Skip this argument to capture frames from a camera.}"
|
||||
"{ imglist | | Pass this flag if image list (i.e. .xml or .yaml) file is passed}"
|
||||
"{ crop | false | Preprocess input image by center cropping.}"
|
||||
//"{ labels | | Path to the text file with labels for detected objects.}"
|
||||
"{ model | | Path to the 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 (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;
|
||||
|
||||
vector<string> classes;
|
||||
static bool readStringList( const string& filename, vector<string>& l )
|
||||
{
|
||||
l.resize(0);
|
||||
FileStorage fs(filename, FileStorage::READ);
|
||||
if( !fs.isOpened() )
|
||||
return false;
|
||||
size_t dir_pos = filename.rfind('/');
|
||||
if (dir_pos == string::npos)
|
||||
dir_pos = filename.rfind('\\');
|
||||
FileNode n = fs.getFirstTopLevelNode();
|
||||
if( n.type() != FileNode::SEQ )
|
||||
return false;
|
||||
FileNodeIterator it = n.begin(), it_end = n.end();
|
||||
for( ; it != it_end; ++it )
|
||||
{
|
||||
string fname = (string)*it;
|
||||
if (dir_pos != string::npos)
|
||||
{
|
||||
string fpath = samples::findFile(filename.substr(0, dir_pos + 1) + fname, false);
|
||||
if (fpath.empty())
|
||||
{
|
||||
fpath = samples::findFile(fname);
|
||||
}
|
||||
fname = fpath;
|
||||
}
|
||||
else
|
||||
{
|
||||
fname = samples::findFile(fname);
|
||||
}
|
||||
l.push_back(fname);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
utils::logging::setLogLevel(utils::logging::LOG_LEVEL_INFO);
|
||||
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
|
||||
if (!parser.has("@alias") || parser.has("help"))
|
||||
{
|
||||
cout << about << endl;
|
||||
parser.printMessage();
|
||||
return -1;
|
||||
}
|
||||
const string modelName = parser.get<String>("@alias");
|
||||
const string zooFile = findFile(parser.get<String>("zoo"));
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
parser.about(about);
|
||||
if (argc == 1 || parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
float scale = parser.get<float>("scale");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
Scalar std = parser.get<Scalar>("std");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
bool crop = parser.get<bool>("crop");
|
||||
int inpWidth = parser.get<int>("width");
|
||||
int inpHeight = parser.get<int>("height");
|
||||
String model = findModel(parser.get<String>("model"), sha1);
|
||||
String backend = parser.get<String>("backend");
|
||||
String target = parser.get<String>("target");
|
||||
bool isImgList = parser.has("imglist");
|
||||
|
||||
// Open file with labels.
|
||||
string labels_filename = parser.get<String>("labels");
|
||||
string file = findFile(labels_filename);
|
||||
ifstream ifs(file.c_str());
|
||||
if (!ifs.is_open()){
|
||||
cout<<"File " << file << " not found";
|
||||
exit(1);
|
||||
}
|
||||
string line;
|
||||
while (getline(ifs, line))
|
||||
{
|
||||
classes.push_back(line);
|
||||
}
|
||||
if (!parser.check())
|
||||
{
|
||||
parser.printErrors();
|
||||
return 1;
|
||||
}
|
||||
CV_Assert(!model.empty());
|
||||
//! [Read and initialize network]
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if (backend != "default" || target != "cpu"){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
Net net = readNetFromONNX(model, engine);
|
||||
net.setPreferableBackend(getBackendID(backend));
|
||||
net.setPreferableTarget(getTargetID(target));
|
||||
net.setProfilingMode(DNN_PROFILE_SUMMARY);
|
||||
//! [Read and initialize network]
|
||||
|
||||
// Create a window
|
||||
static const std::string kWinName = "Deep learning image classification in OpenCV";
|
||||
namedWindow(kWinName, WINDOW_NORMAL);
|
||||
|
||||
//Create FontFace for putText
|
||||
FontFace sans("sans");
|
||||
|
||||
//! [Open a video file or an image file or a camera stream]
|
||||
VideoCapture cap;
|
||||
vector<string> imageList;
|
||||
size_t currentImageIndex = 0;
|
||||
|
||||
if (parser.has("input")) {
|
||||
string input = findFile(parser.get<String>("input"));
|
||||
|
||||
if (isImgList) {
|
||||
bool check = readStringList(samples::findFile(input), imageList);
|
||||
if (imageList.empty() || !check) {
|
||||
cout << "Error: No images found or the provided file is not a valid .yaml or .xml file." << endl;
|
||||
return -1;
|
||||
}
|
||||
} else {
|
||||
// Input is not a directory, try to open as video or image
|
||||
cap.open(input);
|
||||
if (!cap.isOpened()) {
|
||||
cout << "Failed to open the input." << endl;
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
cap.open(0); // Open default camera
|
||||
}
|
||||
//! [Open a video file or an image file or a camera stream]
|
||||
|
||||
Mat frame, blob;
|
||||
for(;;)
|
||||
{
|
||||
if (!imageList.empty()) {
|
||||
// Handling directory of images
|
||||
if (currentImageIndex >= imageList.size()) {
|
||||
waitKey();
|
||||
break; // Exit if all images are processed
|
||||
}
|
||||
frame = imread(imageList[currentImageIndex++]);
|
||||
if(frame.empty()){
|
||||
cout<<"Cannot open file"<<endl;
|
||||
continue;
|
||||
}
|
||||
} else {
|
||||
// Handling video or single image
|
||||
cap >> frame;
|
||||
}
|
||||
if (frame.empty())
|
||||
{
|
||||
break;
|
||||
}
|
||||
//! [Create a 4D blob from a frame]
|
||||
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, crop);
|
||||
// Check std values.
|
||||
if (std.val[0] != 0.0 && std.val[1] != 0.0 && std.val[2] != 0.0)
|
||||
{
|
||||
// Divide blob by std.
|
||||
divide(blob, std, blob);
|
||||
}
|
||||
//! [Create a 4D blob from a frame]
|
||||
//! [Set input blob]
|
||||
net.setInput(blob);
|
||||
//! [Set input blob]
|
||||
|
||||
TickMeter timeRecorder;
|
||||
timeRecorder.reset();
|
||||
Mat prob = net.forward();
|
||||
double t1;
|
||||
//! [Make forward pass]
|
||||
timeRecorder.start();
|
||||
prob = net.forward();
|
||||
timeRecorder.stop();
|
||||
net.printPerfProfile();
|
||||
//! [Make forward pass]
|
||||
|
||||
//! [Get a class with a highest score]
|
||||
int N = (int)prob.total(), K = std::min(5, N);
|
||||
std::vector<std::pair<float, int> > prob_vec;
|
||||
for (int i = 0; i < N; i++) {
|
||||
prob_vec.push_back(std::make_pair(-prob.at<float>(i), i));
|
||||
}
|
||||
std::sort(prob_vec.begin(), prob_vec.end());
|
||||
|
||||
//! [Get a class with a highest score]
|
||||
t1 = timeRecorder.getTimeMilli();
|
||||
timeRecorder.reset();
|
||||
string label = format("Inference time: %.1f ms", t1);
|
||||
Mat subframe = frame(Rect(0, 0, std::min(1000, frame.cols), std::min(300, frame.rows)));
|
||||
subframe *= 0.3f;
|
||||
putText(frame, label, Point(20, 50), Scalar(0, 255, 0), sans, 25, 800);
|
||||
|
||||
// Print predicted class.
|
||||
for (int i = 0; i < K; i++) {
|
||||
int classId = prob_vec[i].second;
|
||||
float confidence = -prob_vec[i].first;
|
||||
label = format("%d. %s: %.2f", i+1, (classes.empty() ? format("Class #%d", classId).c_str() :
|
||||
classes[classId].c_str()), confidence);
|
||||
putText(frame, label, Point(20, 110 + i*35), Scalar(0, 255, 0), sans, 25, 500);
|
||||
}
|
||||
imshow(kWinName, frame);
|
||||
int key = waitKey(isImgList ? 1000 : 100);
|
||||
if (key == ' ')
|
||||
key = waitKey();
|
||||
if (key == 'q' || key == 27) // Check if 'q' or 'ESC' is pressed
|
||||
return 0;
|
||||
}
|
||||
waitKey();
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
import os
|
||||
import glob
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
import sys
|
||||
from common import *
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Firstly, download required models 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.\n"\n
|
||||
|
||||
To run:
|
||||
python classification.py model_name --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)
|
||||
|
||||
Sample command:
|
||||
python classification.py googlenet --input=path/to/image
|
||||
Model path can also be specified using --model argument
|
||||
'''
|
||||
)
|
||||
|
||||
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',
|
||||
help='Path to input image or video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--crop', type=bool, default=False,
|
||||
help='Center crop the image.')
|
||||
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, 'classification')
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Use this script to run classification deep learning networks using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args(func_args)
|
||||
|
||||
def load_images(directory):
|
||||
# List all common image file extensions, feel free to add more if needed
|
||||
extensions = ['jpg', 'jpeg', 'png', 'bmp', 'tif', 'tiff']
|
||||
files = []
|
||||
for extension in extensions:
|
||||
files.extend(glob.glob(os.path.join(directory, f'*.{extension}')))
|
||||
return files
|
||||
|
||||
def main(func_args=None):
|
||||
args = get_args_parser(func_args)
|
||||
if args.alias is None or hasattr(args, 'help'):
|
||||
help()
|
||||
exit(1)
|
||||
|
||||
cv.utils.logging.setLogLevel(cv.utils.logging.LOG_LEVEL_INFO)
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
args.labels = findFile(args.labels)
|
||||
|
||||
# Load names of classes
|
||||
labels = None
|
||||
if args.labels:
|
||||
with open(args.labels, 'rt') as f:
|
||||
labels = f.read().rstrip('\n').split('\n')
|
||||
|
||||
# Load a network
|
||||
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))
|
||||
if hasattr(cv.dnn, 'DNN_PROFILE_SUMMARY'):
|
||||
net.setProfilingMode(cv.dnn.DNN_PROFILE_SUMMARY)
|
||||
|
||||
winName = 'Deep learning image classification in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_NORMAL)
|
||||
|
||||
isdir = False
|
||||
|
||||
if args.input:
|
||||
input_path = args.input
|
||||
|
||||
if os.path.isdir(input_path):
|
||||
isdir = True
|
||||
image_files = load_images(input_path)
|
||||
if not image_files:
|
||||
print("No images found in the directory.")
|
||||
exit(-1)
|
||||
current_image_index = 0
|
||||
else:
|
||||
input_path = findFile(input_path)
|
||||
cap = cv.VideoCapture(input_path)
|
||||
if not cap.isOpened():
|
||||
print("Failed to open the input video")
|
||||
exit(-1)
|
||||
else:
|
||||
cap = cv.VideoCapture(0)
|
||||
|
||||
while cv.waitKey(1) < 0:
|
||||
if isdir:
|
||||
if current_image_index >= len(image_files):
|
||||
break
|
||||
frame = cv.imread(image_files[current_image_index])
|
||||
current_image_index += 1
|
||||
else:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
cv.waitKey()
|
||||
break
|
||||
|
||||
# Create a 4D blob from a frame.
|
||||
inpWidth = args.width if args.width else frame.shape[1]
|
||||
inpHeight = args.height if args.height else frame.shape[0]
|
||||
|
||||
blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=args.crop)
|
||||
if args.std:
|
||||
blob[0] /= np.asarray(args.std, dtype=np.float32).reshape(3, 1, 1)
|
||||
|
||||
# Run a model
|
||||
net.setInput(blob)
|
||||
t0 = cv.getTickCount()
|
||||
out = net.forward()
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
net.printPerfProfile()
|
||||
|
||||
(h, w, _) = frame.shape
|
||||
roi_rows = min(300, h)
|
||||
roi_cols = min(1000, w)
|
||||
frame[:roi_rows,:roi_cols,:] >>= 1
|
||||
|
||||
# Put efficiency information.
|
||||
label = 'Inference time: %.1f ms' % (t * 1000.0)
|
||||
cv.putText(frame, label, (15, 30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0))
|
||||
|
||||
# Print predicted classes.
|
||||
out = out.flatten()
|
||||
K = 5
|
||||
topKidx = np.argpartition(out, -K)[-K:]
|
||||
for i in range(K):
|
||||
classId = topKidx[i]
|
||||
confidence = out[classId]
|
||||
label = '%s: %.2f' % (labels[classId] if labels else 'Class #%d' % classId, confidence)
|
||||
cv.putText(frame, label, (15, 90 + i*30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0))
|
||||
|
||||
cv.imshow(winName, frame)
|
||||
key = cv.waitKey(1000 if isdir else 100)
|
||||
|
||||
if key >= 0:
|
||||
key &= 255
|
||||
if key == ord(' '):
|
||||
key = cv.waitKey() & 255
|
||||
if key == ord('q') or key == 27: # Wait for 1 second on each image, press 'q' to exit
|
||||
sys.exit(0)
|
||||
cv.waitKey()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,121 @@
|
||||
// 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
|
||||
// To download the onnx model, see: https://storage.googleapis.com/ailia-models/colorization/colorizer.onnx
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include "common.hpp"
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <iostream>
|
||||
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
using namespace cv::dnn;
|
||||
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
const string about =
|
||||
"This sample demonstrates recoloring grayscale images with dnn.\n"
|
||||
"This program is based on:\n"
|
||||
" http://richzhang.github.io/colorization\n"
|
||||
" https://github.com/richzhang/colorization\n"
|
||||
"To download the onnx model:\n"
|
||||
" https://storage.googleapis.com/ailia-models/colorization/colorizer.onnx\n";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ input i | baboon.jpg | Path to the input image }"
|
||||
"{ onnx_model_path | | Path to the ONNX model. Required. }";
|
||||
|
||||
const string backend_keys = format(
|
||||
"{ backend | 0 | Choose one of computation backends: "
|
||||
"%d: automatically (by default), "
|
||||
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"%d: OpenCV implementation, "
|
||||
"%d: VKCOM, "
|
||||
"%d: CUDA, "
|
||||
"%d: WebNN }",
|
||||
cv::dnn::DNN_BACKEND_DEFAULT, cv::dnn::DNN_BACKEND_INFERENCE_ENGINE, cv::dnn::DNN_BACKEND_OPENCV,
|
||||
cv::dnn::DNN_BACKEND_VKCOM, cv::dnn::DNN_BACKEND_CUDA, cv::dnn::DNN_BACKEND_WEBNN);
|
||||
const string target_keys = format(
|
||||
"{ target | 0 | Choose one of target computation devices: "
|
||||
"%d: CPU target (by default), "
|
||||
"%d: OpenCL, "
|
||||
"%d: OpenCL fp16 (half-float precision), "
|
||||
"%d: VPU, "
|
||||
"%d: Vulkan, "
|
||||
"%d: CUDA, "
|
||||
"%d: CUDA fp16 (half-float preprocess) }",
|
||||
cv::dnn::DNN_TARGET_CPU, cv::dnn::DNN_TARGET_OPENCL, cv::dnn::DNN_TARGET_OPENCL_FP16,
|
||||
cv::dnn::DNN_TARGET_MYRIAD, cv::dnn::DNN_TARGET_VULKAN, cv::dnn::DNN_TARGET_CUDA,
|
||||
cv::dnn::DNN_TARGET_CUDA_FP16);
|
||||
|
||||
const string keys = param_keys + backend_keys + target_keys;
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
parser.about(about);
|
||||
|
||||
if (parser.has("help")) {
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
string inputImagePath = parser.get<string>("input");
|
||||
string onnxModelPath = parser.get<string>("onnx_model_path");
|
||||
int backendId = parser.get<int>("backend");
|
||||
int targetId = parser.get<int>("target");
|
||||
|
||||
if (onnxModelPath.empty()) {
|
||||
cerr << "The path to the ONNX model is required!" << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
Mat imgGray = imread(samples::findFile(inputImagePath), IMREAD_GRAYSCALE);
|
||||
if (imgGray.empty()) {
|
||||
cerr << "Could not read the image: " << inputImagePath << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
Mat imgL = imgGray;
|
||||
imgL.convertTo(imgL, CV_32F, 100.0/255.0);
|
||||
Mat imgLResized;
|
||||
resize(imgL, imgLResized, Size(256, 256), 0, 0, INTER_CUBIC);
|
||||
|
||||
// Prepare the model
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if (backendId != 0 || targetId != 0){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
dnn::Net net = dnn::readNetFromONNX(onnxModelPath, engine);
|
||||
net.setPreferableBackend(backendId);
|
||||
net.setPreferableTarget(targetId);
|
||||
//! [Read and initialize network]
|
||||
|
||||
// Create blob from the image
|
||||
Mat blob = dnn::blobFromImage(imgLResized, 1.0, Size(256, 256), Scalar(), false, false);
|
||||
|
||||
net.setInput(blob);
|
||||
|
||||
// Run inference
|
||||
Mat result = net.forward();
|
||||
Size siz(result.size[2], result.size[3]);
|
||||
Mat a(siz, CV_32F, result.ptr(0,0));
|
||||
Mat b(siz, CV_32F, result.ptr(0,1));
|
||||
resize(a, a, imgGray.size());
|
||||
resize(b, b, imgGray.size());
|
||||
|
||||
// merge, and convert back to BGR
|
||||
Mat color, chn[] = {imgL, a, b};
|
||||
|
||||
// Proc
|
||||
Mat lab;
|
||||
merge(chn, 3, lab);
|
||||
cvtColor(lab, color, COLOR_Lab2BGR);
|
||||
|
||||
imshow("input image", imgGray);
|
||||
imshow("output image", color);
|
||||
waitKey();
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
# Script is based on https://github.com/richzhang/colorization/blob/master/colorization/colorize.py
|
||||
# To download the onnx model, see: https://storage.googleapis.com/ailia-models/colorization/colorizer.onnx
|
||||
# python colorization.py --onnx_model_path colorizer.onnx --input ansel_adams3.jpg
|
||||
import numpy as np
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
|
||||
def parse_args():
|
||||
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE,
|
||||
cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
|
||||
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD,
|
||||
cv.dnn.DNN_TARGET_HDDL, cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)
|
||||
|
||||
parser = argparse.ArgumentParser(description='iColor: deep interactive colorization')
|
||||
parser.add_argument('--input', default='baboon.jpg',help='Path to image.')
|
||||
parser.add_argument('--onnx_model_path', help='Path to onnx model', required=True)
|
||||
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
|
||||
help="Choose one of computation backends: "
|
||||
"%d: automatically (by default), "
|
||||
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"%d: OpenCV implementation, "
|
||||
"%d: VKCOM, "
|
||||
"%d: CUDA" % backends)
|
||||
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
|
||||
help='Choose one of target computation devices: '
|
||||
'%d: CPU target (by default), '
|
||||
'%d: OpenCL, '
|
||||
'%d: OpenCL fp16 (half-float precision), '
|
||||
'%d: NCS2 VPU, '
|
||||
'%d: HDDL VPU, '
|
||||
'%d: Vulkan, '
|
||||
'%d: CUDA, '
|
||||
'%d: CUDA fp16 (half-float preprocess)'% targets)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
img_gray=cv.imread(cv.samples.findFile(args.input),cv.IMREAD_GRAYSCALE)
|
||||
|
||||
img_gray_rs = cv.resize(img_gray, (256, 256), interpolation=cv.INTER_CUBIC)
|
||||
img_gray_rs = img_gray_rs.astype(np.float32) # Convert to float to avoid data overflow
|
||||
img_gray_rs *= (100.0 / 255.0) # Scale L channel to 0-100 range
|
||||
|
||||
onnx_model_path = args.onnx_model_path # Update this path to your ONNX model's path
|
||||
engine = cv.dnn.ENGINE_AUTO
|
||||
if args.backend != 0 or args.target != 0:
|
||||
engine = cv.dnn.ENGINE_CLASSIC
|
||||
session = cv.dnn.readNetFromONNX(onnx_model_path, engine)
|
||||
session.setPreferableBackend(args.backend)
|
||||
session.setPreferableTarget(args.target)
|
||||
|
||||
# Process each image in the batch (assuming batch processing is needed)
|
||||
blob = cv.dnn.blobFromImage(img_gray_rs, swapRB=False) # Adjust swapRB according to your model's training
|
||||
session.setInput(blob)
|
||||
result_numpy = np.array(session.forward()[0])
|
||||
|
||||
if result_numpy.shape[0] == 2:
|
||||
# Transpose result_numpy to shape (H, W, 2)
|
||||
ab = result_numpy.transpose((1, 2, 0))
|
||||
else:
|
||||
# If it's already (H, W, 2), assign it directly
|
||||
ab = result_numpy
|
||||
|
||||
|
||||
# Resize ab to match img_gray's dimensions if they are not the same
|
||||
h, w = img_gray.shape
|
||||
if ab.shape[:2] != (h, w):
|
||||
ab_resized = cv.resize(ab, (w, h), interpolation=cv.INTER_LINEAR)
|
||||
else:
|
||||
ab_resized = ab
|
||||
|
||||
# Expand dimensions of L to match ab's dimensions
|
||||
img_l_expanded = np.expand_dims(img_gray, axis=-1)
|
||||
|
||||
# Concatenate L with AB to get the LAB image
|
||||
lab_image = np.concatenate((img_l_expanded, ab_resized), axis=-1)
|
||||
|
||||
# Convert the Lab image to a 32-bit float format
|
||||
lab_image = lab_image.astype(np.float32)
|
||||
|
||||
# Normalize L channel to the range [0, 100] and AB channels to the range [-127, 127]
|
||||
lab_image[:, :, 0] *= (100.0 / 255.0) # Rescale L channel
|
||||
#lab_image[:, :, 1:] -= 128 # Shift AB channels
|
||||
|
||||
# Convert the LAB image to BGR
|
||||
image_bgr_out = cv.cvtColor(lab_image, cv.COLOR_Lab2BGR)
|
||||
cv.imshow("input image",img_gray)
|
||||
cv.imshow("output image",image_bgr_out)
|
||||
cv.waitKey(0)
|
||||
@@ -0,0 +1,218 @@
|
||||
#include <opencv2/core/utils/filesystem.hpp>
|
||||
#include<iostream>
|
||||
using namespace cv;
|
||||
|
||||
std::string genArgument(const std::string& argName, const std::string& help,
|
||||
const std::string& modelName, const std::string& zooFile,
|
||||
char key = ' ', std::string defaultVal = "");
|
||||
|
||||
std::string genPreprocArguments(const std::string& modelName, const std::string& zooFile, const std::string& prefix);
|
||||
|
||||
std::string findFile(const std::string& filename);
|
||||
|
||||
std::string findModel(const std::string& filename, const std::string& sha1);
|
||||
|
||||
std::vector<std::string> findAliases(std::string& zooFile, const std::string& sampleType);
|
||||
|
||||
inline int getBackendID(const String& backend) {
|
||||
std::map<String, int> backendIDs = {
|
||||
{"default", cv::dnn::DNN_BACKEND_DEFAULT},
|
||||
{"openvino", cv::dnn::DNN_BACKEND_INFERENCE_ENGINE},
|
||||
{"opencv", cv::dnn::DNN_BACKEND_OPENCV},
|
||||
{"vkcom", cv::dnn::DNN_BACKEND_VKCOM},
|
||||
{"cuda", cv::dnn::DNN_BACKEND_CUDA},
|
||||
{"webnn", cv::dnn::DNN_BACKEND_WEBNN}
|
||||
};
|
||||
if(backendIDs.find(backend) != backendIDs.end()){
|
||||
return backendIDs[backend];
|
||||
}else {
|
||||
throw std::invalid_argument("Invalid backend name: " + backend);
|
||||
}
|
||||
}
|
||||
|
||||
inline int getTargetID(const String& target) {
|
||||
std::map<String, int> targetIDs = {
|
||||
{"cpu", cv::dnn::DNN_TARGET_CPU},
|
||||
{"opencl", cv::dnn::DNN_TARGET_OPENCL},
|
||||
{"opencl_fp16", cv::dnn::DNN_TARGET_OPENCL_FP16},
|
||||
{"vpu", cv::dnn::DNN_TARGET_MYRIAD},
|
||||
{"vulkan", cv::dnn::DNN_TARGET_VULKAN},
|
||||
{"cuda", cv::dnn::DNN_TARGET_CUDA},
|
||||
{"cuda_fp16", cv::dnn::DNN_TARGET_CUDA_FP16}
|
||||
};
|
||||
if(targetIDs.find(target) != targetIDs.end()){
|
||||
return targetIDs[target];
|
||||
}else {
|
||||
throw std::invalid_argument("Invalid target name: " + target);
|
||||
}
|
||||
}
|
||||
|
||||
std::string genArgument(const std::string& argName, const std::string& help,
|
||||
const std::string& modelName, const std::string& zooFile,
|
||||
char key, std::string defaultVal)
|
||||
{
|
||||
if (!modelName.empty())
|
||||
{
|
||||
FileStorage fs(zooFile, FileStorage::READ);
|
||||
if (fs.isOpened())
|
||||
{
|
||||
FileNode node = fs[modelName];
|
||||
if (!node.empty())
|
||||
{
|
||||
FileNode value = node[argName];
|
||||
if (argName.find("sha1") != std::string::npos) {
|
||||
std::string prefix = argName.substr(0, argName.find("sha1"));
|
||||
if (prefix == "config_"){
|
||||
value = node[prefix+"load_info"]["sha1"];
|
||||
}
|
||||
else{
|
||||
value = node[prefix+"load_info"][argName];
|
||||
}
|
||||
}
|
||||
if (argName.find("download_sha") != std::string::npos) {
|
||||
std::string prefix = argName.substr(0, argName.find("download_sha"));
|
||||
value = node[prefix+"load_info"][argName];
|
||||
}
|
||||
if (!value.empty())
|
||||
{
|
||||
if (value.isReal())
|
||||
defaultVal = format("%f", (float)value);
|
||||
else if (value.isString())
|
||||
defaultVal = (std::string)value;
|
||||
else if (value.isInt())
|
||||
defaultVal = format("%d", (int)value);
|
||||
else if (value.isSeq())
|
||||
{
|
||||
for (size_t i = 0; i < value.size(); ++i)
|
||||
{
|
||||
FileNode v = value[(int)i];
|
||||
if (v.isInt())
|
||||
defaultVal += format("%d ", (int)v);
|
||||
else if (v.isReal())
|
||||
defaultVal += format("%f ", (float)v);
|
||||
else
|
||||
CV_Error(Error::StsNotImplemented, "Unexpected value format");
|
||||
}
|
||||
}
|
||||
else
|
||||
CV_Error(Error::StsNotImplemented, "Unexpected field format");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return "{ " + argName + " " + key + " | " + defaultVal + " | " + help + " }";
|
||||
}
|
||||
|
||||
std::string findModel(const std::string& filename, const std::string& sha1)
|
||||
{
|
||||
if (filename.empty() || utils::fs::exists(filename))
|
||||
return filename;
|
||||
|
||||
if(!getenv("OPENCV_DOWNLOAD_CACHE_DIR")){
|
||||
std::cout<< "[WARN] Please specify a path to model download directory in OPENCV_DOWNLOAD_CACHE_DIR environment variable"<<std::endl;
|
||||
return findFile(filename);
|
||||
}
|
||||
else{
|
||||
std::string modelPath = utils::fs::join(getenv("OPENCV_DOWNLOAD_CACHE_DIR"), utils::fs::join(sha1, filename));
|
||||
if (utils::fs::exists(modelPath))
|
||||
return modelPath;
|
||||
modelPath = utils::fs::join(getenv("OPENCV_DOWNLOAD_CACHE_DIR"),filename);
|
||||
if (utils::fs::exists(modelPath))
|
||||
return modelPath;
|
||||
}
|
||||
|
||||
std::cout << "File " + filename + " not found! "
|
||||
<< "Please specify a path to model download directory in OPENCV_DOWNLOAD_CACHE_DIR "
|
||||
<< "environment variable or pass a full path to " + filename
|
||||
<< std::endl;
|
||||
std::exit(1);
|
||||
}
|
||||
|
||||
std::string findFile(const std::string& filename)
|
||||
{
|
||||
if (filename.empty() || utils::fs::exists(filename))
|
||||
return filename;
|
||||
|
||||
if(!getenv("OPENCV_SAMPLES_DATA_PATH")){
|
||||
std::cout<< "[WARN] Please specify a path to opencv/samples/data in OPENCV_SAMPLES_DATA_PATH environment variable"<<std::endl;
|
||||
}
|
||||
else{
|
||||
std::string samplePath = utils::fs::join(getenv("OPENCV_SAMPLES_DATA_PATH"), filename);
|
||||
if (utils::fs::exists(samplePath))
|
||||
return samplePath;
|
||||
}
|
||||
const char* extraPaths[] = {getenv("OPENCV_SAMPLES_DATA_PATH"),
|
||||
getenv("OPENCV_DNN_TEST_DATA_PATH"),
|
||||
getenv("OPENCV_TEST_DATA_PATH")};
|
||||
for (int i = 0; i < 3; ++i)
|
||||
{
|
||||
if (extraPaths[i] == NULL)
|
||||
continue;
|
||||
std::string absPath = utils::fs::join(extraPaths[i], utils::fs::join("dnn", filename));
|
||||
if (utils::fs::exists(absPath))
|
||||
return absPath;
|
||||
}
|
||||
std::cout << "File " + filename + " not found! "
|
||||
<< "Please specify the path to /opencv/samples/data in the OPENCV_SAMPLES_DATA_PATH environment variable, "
|
||||
<< "or specify the path to opencv_extra/testdata in the OPENCV_DNN_TEST_DATA_PATH environment variable, "
|
||||
<< "or specify the path to the model download cache directory in the OPENCV_DOWNLOAD_CACHE_DIR environment variable, "
|
||||
<< "or pass the full path to " + filename + "."
|
||||
<< std::endl;
|
||||
std::exit(1);
|
||||
}
|
||||
|
||||
std::string genPreprocArguments(const std::string& modelName, const std::string& zooFile, const std::string& prefix="")
|
||||
{
|
||||
return genArgument(prefix + "model", "Path to a binary file of model contains trained weights. "
|
||||
"It could be a file with extensions .caffemodel (Caffe), "
|
||||
".pb (TensorFlow), .bin (OpenVINO).",
|
||||
modelName, zooFile, 'm') +
|
||||
genArgument(prefix + "config", "Path to a text file of model contains network configuration. "
|
||||
"It could be a file with extensions .prototxt (Caffe), .pbtxt (TensorFlow), .xml (OpenVINO).",
|
||||
modelName, zooFile, 'c') +
|
||||
genArgument(prefix + "mean", "Preprocess input image by subtracting mean values. Mean values should be in BGR order and delimited by spaces.",
|
||||
modelName, zooFile) +
|
||||
genArgument(prefix + "std", "Preprocess input image by dividing on a standard deviation.",
|
||||
modelName, zooFile) +
|
||||
genArgument(prefix + "scale", "Preprocess input image by multiplying on a scale factor.",
|
||||
modelName, zooFile, ' ', "1.0") +
|
||||
genArgument(prefix + "width", "Preprocess input image by resizing to a specific width.",
|
||||
modelName, zooFile, ' ', "-1") +
|
||||
genArgument(prefix + "height", "Preprocess input image by resizing to a specific height.",
|
||||
modelName, zooFile, ' ', "-1") +
|
||||
genArgument(prefix + "rgb", "Indicate that model works with RGB input images instead BGR ones.",
|
||||
modelName, zooFile)+
|
||||
genArgument(prefix + "labels", "Path to a text file with names of classes to label detected objects.",
|
||||
modelName, zooFile)+
|
||||
genArgument(prefix + "postprocessing", "Indicate the postprocessing type of model i.e. yolov8, yolonas, etc.",
|
||||
modelName, zooFile)+
|
||||
genArgument(prefix + "sha1", "Optional path to hashsum of downloaded model to be loaded from models.yml",
|
||||
modelName, zooFile)+
|
||||
genArgument(prefix + "config_sha1", "Optional path to hashsum of downloaded config to be loaded from models.yml",
|
||||
modelName, zooFile)+
|
||||
genArgument(prefix + "download_sha", "Optional path to hashsum of downloaded model to be loaded from models.yml",
|
||||
modelName, zooFile);
|
||||
}
|
||||
|
||||
std::vector<std::string> findAliases(std::string& zooFile, const std::string& sampleType) {
|
||||
std::vector<std::string> aliases;
|
||||
|
||||
zooFile = findFile(zooFile);
|
||||
|
||||
cv::FileStorage fs(zooFile, cv::FileStorage::READ);
|
||||
|
||||
cv::FileNode root = fs.root();
|
||||
for (const auto& node : root) {
|
||||
std::string alias = node.name();
|
||||
cv::FileNode sampleNode = node["sample"];
|
||||
|
||||
if (!sampleNode.empty() && sampleNode.isString()) {
|
||||
std::string sampleValue = (std::string)sampleNode;
|
||||
if (sampleValue == sampleType) {
|
||||
aliases.push_back(alias);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return aliases;
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
import sys
|
||||
import os
|
||||
import cv2 as cv
|
||||
|
||||
|
||||
def add_argument(zoo, parser, name, help, required=False, default=None, type=None, action=None, nargs=None, alias=None):
|
||||
if alias is not None:
|
||||
modelName = alias
|
||||
elif len(sys.argv) > 1:
|
||||
modelName = sys.argv[1]
|
||||
else:
|
||||
return
|
||||
|
||||
if os.path.isfile(zoo):
|
||||
fs = cv.FileStorage(zoo, cv.FILE_STORAGE_READ)
|
||||
node = fs.getNode(modelName)
|
||||
if not node.empty():
|
||||
value = node.getNode(name)
|
||||
if "sha1" in name:
|
||||
prefix = name.replace("sha1", "")
|
||||
value = node.getNode(prefix + "load_info")
|
||||
if prefix == "config_":
|
||||
value = value.getNode("sha1")
|
||||
else:
|
||||
value = value.getNode(name)
|
||||
if "download_sha" in name:
|
||||
prefix = name.replace("download_sha", "")
|
||||
value = node.getNode(prefix + "load_info")
|
||||
value = value.getNode(name)
|
||||
if not value.empty():
|
||||
if value.isReal():
|
||||
default = value.real()
|
||||
elif value.isString():
|
||||
default = value.string()
|
||||
elif value.isInt():
|
||||
default = int(value.real())
|
||||
elif value.isSeq():
|
||||
default = []
|
||||
for i in range(value.size()):
|
||||
v = value.at(i)
|
||||
if v.isInt():
|
||||
default.append(int(v.real()))
|
||||
elif v.isReal():
|
||||
default.append(v.real())
|
||||
else:
|
||||
print('Unexpected value format')
|
||||
exit(0)
|
||||
else:
|
||||
print('Unexpected field format')
|
||||
exit(0)
|
||||
required = False
|
||||
|
||||
if action == 'store_true':
|
||||
default = 1 if default == 'true' else (0 if default == 'false' else default)
|
||||
assert(default is None or default == 0 or default == 1)
|
||||
parser.add_argument('--' + name, required=required, help=help, default=bool(default),
|
||||
action=action)
|
||||
else:
|
||||
parser.add_argument('--' + name, required=required, help=help, default=default,
|
||||
action=action, nargs=nargs, type=type)
|
||||
|
||||
|
||||
def add_preproc_args(zoo, parser, sample, alias=None, prefix=""):
|
||||
aliases = []
|
||||
if os.path.isfile(zoo) and prefix == "":
|
||||
fs = cv.FileStorage(zoo, cv.FILE_STORAGE_READ)
|
||||
root = fs.root()
|
||||
for name in root.keys():
|
||||
model = root.getNode(name)
|
||||
if model.getNode('sample').string() == sample:
|
||||
aliases.append(name)
|
||||
if len(aliases):
|
||||
parser.add_argument(prefix+'alias', nargs='?', choices=aliases,
|
||||
help='An alias name of model to extract preprocessing parameters from models.yml file.')
|
||||
|
||||
add_argument(zoo, parser, prefix+'model',
|
||||
help='Path to a binary file of model contains trained weights. '
|
||||
'It could be a file with extensions .caffemodel (Caffe), '
|
||||
'.pb (TensorFlow), .bin (OpenVINO)', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'config',
|
||||
help='Path to a text file of model contains network configuration. '
|
||||
'It could be a file with extensions .prototxt (Caffe), .pbtxt or .config (TensorFlow), .xml (OpenVINO)', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'mean', nargs='+', type=float, default=[0, 0, 0],
|
||||
help='Preprocess input image by subtracting mean values. '
|
||||
'Mean values should be in BGR order.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'std', nargs='+', type=float, default=[0, 0, 0],
|
||||
help='Preprocess input image by dividing on a standard deviation.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'scale', type=float, default=1.0,
|
||||
help='Preprocess input image by multiplying on a scale factor.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'width', type=int,
|
||||
help='Preprocess input image by resizing to a specific width.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'height', type=int,
|
||||
help='Preprocess input image by resizing to a specific height.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'rgb', action='store_true',
|
||||
help='Indicate that model works with RGB input images instead BGR ones.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'labels',
|
||||
help='Optional path to a text file with names of labels to label detected objects.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'postprocessing', type=str,
|
||||
help='Post-processing kind depends on model topology.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'background_label_id', type=int, default=-1,
|
||||
help='An index of background class in predictions. If not negative, exclude such class from list of classes.', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'sha1', type=str,
|
||||
help='Optional path to hashsum of downloaded model to be loaded from models.yml', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'config_sha1', type=str,
|
||||
help='Optional path to hashsum of downloaded config to be loaded from models.yml', alias=alias)
|
||||
add_argument(zoo, parser, prefix+'download_sha', type=str,
|
||||
help='Optional path to hashsum of downloaded model to be loaded from models.yml', alias=alias)
|
||||
|
||||
def findModel(filename, sha1):
|
||||
if filename:
|
||||
if os.path.exists(filename):
|
||||
return filename
|
||||
|
||||
fpath = cv.samples.findFile(filename, False)
|
||||
if fpath:
|
||||
return fpath
|
||||
|
||||
if os.getenv('OPENCV_DOWNLOAD_CACHE_DIR') is None:
|
||||
print('[WARN] Please specify a path to model download directory in OPENCV_DOWNLOAD_CACHE_DIR environment variable.')
|
||||
return findFile(filename)
|
||||
|
||||
if os.path.exists(os.path.join(os.environ['OPENCV_DOWNLOAD_CACHE_DIR'], sha1, filename)):
|
||||
return os.path.join(os.environ['OPENCV_DOWNLOAD_CACHE_DIR'], sha1, filename)
|
||||
|
||||
if os.path.exists(os.path.join(os.environ['OPENCV_DOWNLOAD_CACHE_DIR'], filename)):
|
||||
return os.path.join(os.environ['OPENCV_DOWNLOAD_CACHE_DIR'], filename)
|
||||
|
||||
raise FileNotFoundError('File ' + filename + ' not found! Please specify a path to '
|
||||
'model download directory in OPENCV_DOWNLOAD_CACHE_DIR '
|
||||
'environment variable or pass a full path to ' + filename)
|
||||
|
||||
def findFile(filename):
|
||||
if filename:
|
||||
if os.path.exists(filename):
|
||||
return filename
|
||||
|
||||
fpath = cv.samples.findFile(filename, False)
|
||||
if fpath:
|
||||
return fpath
|
||||
|
||||
if os.getenv('OPENCV_SAMPLES_DATA_PATH') is None:
|
||||
print('[WARN] Please specify a path to `/samples/data` in OPENCV_SAMPLES_DATA_PATH environment variable.')
|
||||
exit(0)
|
||||
|
||||
if os.path.exists(os.path.join(os.environ['OPENCV_SAMPLES_DATA_PATH'], filename)):
|
||||
return os.path.join(os.environ['OPENCV_SAMPLES_DATA_PATH'], filename)
|
||||
|
||||
for path in ['OPENCV_DNN_TEST_DATA_PATH', 'OPENCV_TEST_DATA_PATH', 'OPENCV_SAMPLES_DATA_PATH']:
|
||||
try:
|
||||
extraPath = os.environ[path]
|
||||
absPath = os.path.join(extraPath, 'dnn', filename)
|
||||
if os.path.exists(absPath):
|
||||
return absPath
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
raise FileNotFoundError(
|
||||
'File ' + filename + ' not found! Please specify the path to '
|
||||
'/opencv/samples/data in the OPENCV_SAMPLES_DATA_PATH environment variable, '
|
||||
'or specify the path to opencv_extra/testdata in the OPENCV_DNN_TEST_DATA_PATH environment variable, '
|
||||
'or specify the path to the model download cache directory in the OPENCV_DOWNLOAD_CACHE_DIR environment variable, '
|
||||
'or pass the full path to ' + filename + '.'
|
||||
)
|
||||
|
||||
|
||||
def get_backend_id(backend_name):
|
||||
backend_ids = {
|
||||
"default": cv.dnn.DNN_BACKEND_DEFAULT,
|
||||
"openvino": cv.dnn.DNN_BACKEND_INFERENCE_ENGINE,
|
||||
"opencv": cv.dnn.DNN_BACKEND_OPENCV,
|
||||
"vkcom": cv.dnn.DNN_BACKEND_VKCOM,
|
||||
"cuda": cv.dnn.DNN_BACKEND_CUDA
|
||||
}
|
||||
|
||||
if backend_name not in backend_ids:
|
||||
raise ValueError(f"Invalid backend name: {backend_name}")
|
||||
|
||||
return backend_ids[backend_name]
|
||||
|
||||
def get_target_id(target_name):
|
||||
target_ids = {
|
||||
"cpu": cv.dnn.DNN_TARGET_CPU,
|
||||
"opencl": cv.dnn.DNN_TARGET_OPENCL,
|
||||
"opencl_fp16": cv.dnn.DNN_TARGET_OPENCL_FP16,
|
||||
"ncs2_vpu": cv.dnn.DNN_TARGET_MYRIAD,
|
||||
"hddl_vpu": cv.dnn.DNN_TARGET_HDDL,
|
||||
"vulkan": cv.dnn.DNN_TARGET_VULKAN,
|
||||
"cuda": cv.dnn.DNN_TARGET_CUDA,
|
||||
"cuda_fp16": cv.dnn.DNN_TARGET_CUDA_FP16
|
||||
}
|
||||
if target_name not in target_ids:
|
||||
raise ValueError(f"Invalid target name: {target_name}")
|
||||
|
||||
return target_ids[target_name]
|
||||
@@ -0,0 +1,31 @@
|
||||
import cv2 as cv
|
||||
|
||||
#! [CropLayer]
|
||||
class CropLayer(object):
|
||||
def __init__(self, params, blobs):
|
||||
self.xstart = 0
|
||||
self.xend = 0
|
||||
self.ystart = 0
|
||||
self.yend = 0
|
||||
|
||||
# Our layer receives two inputs. We need to crop the first input blob
|
||||
# to match a shape of the second one (keeping batch size and number of channels)
|
||||
def getMemoryShapes(self, inputs):
|
||||
inputShape, targetShape = inputs[0], inputs[1]
|
||||
batchSize, numChannels = inputShape[0], inputShape[1]
|
||||
height, width = targetShape[2], targetShape[3]
|
||||
|
||||
self.ystart = (inputShape[2] - targetShape[2]) // 2
|
||||
self.xstart = (inputShape[3] - targetShape[3]) // 2
|
||||
self.yend = self.ystart + height
|
||||
self.xend = self.xstart + width
|
||||
|
||||
return [[batchSize, numChannels, height, width]]
|
||||
|
||||
def forward(self, inputs):
|
||||
return [inputs[0][:,:,self.ystart:self.yend,self.xstart:self.xend]]
|
||||
#! [CropLayer]
|
||||
|
||||
#! [Register]
|
||||
cv.dnn_registerLayer('Crop', CropLayer)
|
||||
#! [Register]
|
||||
@@ -0,0 +1,121 @@
|
||||
// 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) 2026, BigVision LLC, all rights reserved.
|
||||
// Third party copyrights are property of their respective owners.
|
||||
//
|
||||
// Companion sample for tutorial:
|
||||
// doc/tutorials/dnn/dnn_custom_layers/dnn_custom_layers.md (ONNX section)
|
||||
//
|
||||
// Models used by this sample are generated by:
|
||||
// opencv_extra/testdata/dnn/onnx/generate_custom_layer_models.py
|
||||
// and live at:
|
||||
// opencv_extra/testdata/dnn/onnx/models/custom_layer_default_domain.onnx
|
||||
// opencv_extra/testdata/dnn/onnx/models/custom_layer_custom_domain.onnx
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/core.hpp>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
using namespace cv;
|
||||
using namespace cv::dnn;
|
||||
using namespace std;
|
||||
|
||||
//! [CustomScaleBiasLayer]
|
||||
// y = scale * x + bias, with scale/bias read from ONNX node attributes.
|
||||
class CustomScaleBiasLayer CV_FINAL : public Layer
|
||||
{
|
||||
public:
|
||||
CustomScaleBiasLayer(const LayerParams& params) : Layer(params)
|
||||
{
|
||||
scale = params.get<float>("scale", 1.f);
|
||||
bias = params.get<float>("bias", 0.f);
|
||||
}
|
||||
|
||||
static Ptr<Layer> create(LayerParams& params)
|
||||
{
|
||||
return makePtr<CustomScaleBiasLayer>(params);
|
||||
}
|
||||
|
||||
bool getMemoryShapes(const vector<MatShape>& inpts,
|
||||
const int /*requiredOutputs*/,
|
||||
vector<MatShape>& outShapes,
|
||||
vector<MatShape>& /*internals*/) const CV_OVERRIDE
|
||||
{
|
||||
outShapes.assign(1, inpts[0]);
|
||||
return false;
|
||||
}
|
||||
|
||||
void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr,
|
||||
OutputArrayOfArrays) CV_OVERRIDE
|
||||
{
|
||||
vector<Mat> inps, outs;
|
||||
inputs_arr.getMatVector(inps);
|
||||
outputs_arr.getMatVector(outs);
|
||||
inps[0].convertTo(outs[0], outs[0].type(), scale, bias);
|
||||
}
|
||||
|
||||
private:
|
||||
float scale, bias;
|
||||
};
|
||||
//! [CustomScaleBiasLayer]
|
||||
|
||||
static void printShape(const string& tag, const Mat& m)
|
||||
{
|
||||
cout << tag;
|
||||
for (int d = 0; d < m.dims; ++d)
|
||||
cout << (d ? "x" : "") << m.size[d];
|
||||
cout << "\n";
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
const string keys =
|
||||
"{ help h | | Print help message }"
|
||||
"{ model m | | Path to ONNX model containing the custom op }"
|
||||
"{ op o | MyCustomOp | Op key for registration. Default-domain: just the op_type "
|
||||
"(e.g. MyCustomOp). Custom-domain: <domain>.<op_type> "
|
||||
"(e.g. my.namespace.MyDomainOp) }";
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
parser.about("Demonstrates importing an ONNX model that contains a custom (non-standard) op "
|
||||
"by registering a user-defined layer with cv::dnn::LayerFactory.");
|
||||
if (parser.has("help") || !parser.has("model"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
const string modelPath = parser.get<string>("model");
|
||||
const string opKey = parser.get<string>("op");
|
||||
|
||||
//! [Register CustomScaleBiasLayer]
|
||||
// ONNX op-type lookup: layers in the default `ai.onnx` domain are registered
|
||||
// under their op_type; layers in a non-default domain are registered under
|
||||
// "<domain>.<op_type>" (e.g. "my.namespace.MyDomainOp").
|
||||
LayerFactory::registerLayer(opKey, CustomScaleBiasLayer::create);
|
||||
//! [Register CustomScaleBiasLayer]
|
||||
|
||||
Net net = readNetFromONNX(modelPath);
|
||||
if (net.empty())
|
||||
{
|
||||
cerr << "Failed to load model: " << modelPath << "\n";
|
||||
LayerFactory::unregisterLayer(opKey);
|
||||
return 1;
|
||||
}
|
||||
|
||||
// The companion ONNX models accept a 1x3x4x4 float input.
|
||||
Mat input(vector<int>{1, 3, 4, 4}, CV_32F, Scalar(1.0f));
|
||||
net.setInput(input);
|
||||
Mat out = net.forward();
|
||||
|
||||
cout << "Loaded " << modelPath << " using custom op key '" << opKey << "'.\n";
|
||||
printShape("Input shape: ", input);
|
||||
printShape("Output shape: ", out);
|
||||
cout << "Output[0,0,0,0] = " << out.ptr<float>()[0]
|
||||
<< " (= scale * 1.0 + bias for the registered op)\n";
|
||||
|
||||
LayerFactory::unregisterLayer(opKey);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
#ifndef __OPENCV_SAMPLES_DNN_CUSTOM_LAYERS__
|
||||
#define __OPENCV_SAMPLES_DNN_CUSTOM_LAYERS__
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/dnn/shape_utils.hpp> // getPlane
|
||||
|
||||
//! [InterpLayer]
|
||||
class InterpLayer : public cv::dnn::Layer
|
||||
{
|
||||
public:
|
||||
InterpLayer(const cv::dnn::LayerParams ¶ms) : Layer(params)
|
||||
{
|
||||
outWidth = params.get<int>("width", 0);
|
||||
outHeight = params.get<int>("height", 0);
|
||||
}
|
||||
|
||||
static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params)
|
||||
{
|
||||
return cv::Ptr<cv::dnn::Layer>(new InterpLayer(params));
|
||||
}
|
||||
|
||||
virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
|
||||
const int requiredOutputs,
|
||||
std::vector<std::vector<int> > &outputs,
|
||||
std::vector<std::vector<int> > &internals) const CV_OVERRIDE
|
||||
{
|
||||
CV_UNUSED(requiredOutputs); CV_UNUSED(internals);
|
||||
std::vector<int> outShape(4);
|
||||
outShape[0] = inputs[0][0]; // batch size
|
||||
outShape[1] = inputs[0][1]; // number of channels
|
||||
outShape[2] = outHeight;
|
||||
outShape[3] = outWidth;
|
||||
outputs.assign(1, outShape);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Implementation of this custom layer is based on https://github.com/cdmh/deeplab-public/blob/master/src/caffe/layers/interp_layer.cpp
|
||||
virtual void forward(cv::InputArrayOfArrays inputs_arr,
|
||||
cv::OutputArrayOfArrays outputs_arr,
|
||||
cv::OutputArrayOfArrays internals_arr) CV_OVERRIDE
|
||||
{
|
||||
if (inputs_arr.depth() == CV_16S)
|
||||
{
|
||||
// In case of DNN_TARGET_OPENCL_FP16 target the following method
|
||||
// converts data from FP16 to FP32 and calls this forward again.
|
||||
forward_fallback(inputs_arr, outputs_arr, internals_arr);
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<cv::Mat> inputs, outputs;
|
||||
inputs_arr.getMatVector(inputs);
|
||||
outputs_arr.getMatVector(outputs);
|
||||
|
||||
cv::Mat& inp = inputs[0];
|
||||
cv::Mat& out = outputs[0];
|
||||
const float* inpData = (float*)inp.data;
|
||||
float* outData = (float*)out.data;
|
||||
|
||||
const int batchSize = inp.size[0];
|
||||
const int numChannels = inp.size[1];
|
||||
const int inpHeight = inp.size[2];
|
||||
const int inpWidth = inp.size[3];
|
||||
|
||||
const float rheight = (outHeight > 1) ? static_cast<float>(inpHeight - 1) / (outHeight - 1) : 0.f;
|
||||
const float rwidth = (outWidth > 1) ? static_cast<float>(inpWidth - 1) / (outWidth - 1) : 0.f;
|
||||
for (int h2 = 0; h2 < outHeight; ++h2)
|
||||
{
|
||||
const float h1r = rheight * h2;
|
||||
const int h1 = static_cast<int>(h1r);
|
||||
const int h1p = (h1 < inpHeight - 1) ? 1 : 0;
|
||||
const float h1lambda = h1r - h1;
|
||||
const float h0lambda = 1.f - h1lambda;
|
||||
for (int w2 = 0; w2 < outWidth; ++w2)
|
||||
{
|
||||
const float w1r = rwidth * w2;
|
||||
const int w1 = static_cast<int>(w1r);
|
||||
const int w1p = (w1 < inpWidth - 1) ? 1 : 0;
|
||||
const float w1lambda = w1r - w1;
|
||||
const float w0lambda = 1.f - w1lambda;
|
||||
const float* pos1 = inpData + h1 * inpWidth + w1;
|
||||
float* pos2 = outData + h2 * outWidth + w2;
|
||||
for (int c = 0; c < batchSize * numChannels; ++c)
|
||||
{
|
||||
pos2[0] =
|
||||
h0lambda * (w0lambda * pos1[0] + w1lambda * pos1[w1p]) +
|
||||
h1lambda * (w0lambda * pos1[h1p * inpWidth] + w1lambda * pos1[h1p * inpWidth + w1p]);
|
||||
pos1 += inpWidth * inpHeight;
|
||||
pos2 += outWidth * outHeight;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
int outWidth, outHeight;
|
||||
};
|
||||
//! [InterpLayer]
|
||||
|
||||
//! [ResizeBilinearLayer]
|
||||
class ResizeBilinearLayer CV_FINAL : public cv::dnn::Layer
|
||||
{
|
||||
public:
|
||||
ResizeBilinearLayer(const cv::dnn::LayerParams ¶ms) : Layer(params)
|
||||
{
|
||||
CV_Assert(!params.get<bool>("align_corners", false));
|
||||
CV_Assert(!blobs.empty());
|
||||
|
||||
for (size_t i = 0; i < blobs.size(); ++i)
|
||||
CV_Assert(blobs[i].type() == CV_32SC1);
|
||||
|
||||
// There are two cases of input blob: a single blob which contains output
|
||||
// shape and two blobs with scaling factors.
|
||||
if (blobs.size() == 1)
|
||||
{
|
||||
CV_Assert(blobs[0].total() == 2);
|
||||
outHeight = blobs[0].at<int>(0, 0);
|
||||
outWidth = blobs[0].at<int>(0, 1);
|
||||
factorHeight = factorWidth = 0;
|
||||
}
|
||||
else
|
||||
{
|
||||
CV_Assert(blobs.size() == 2); CV_Assert(blobs[0].total() == 1); CV_Assert(blobs[1].total() == 1);
|
||||
factorHeight = blobs[0].at<int>(0, 0);
|
||||
factorWidth = blobs[1].at<int>(0, 0);
|
||||
outHeight = outWidth = 0;
|
||||
}
|
||||
}
|
||||
|
||||
static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params)
|
||||
{
|
||||
return cv::Ptr<cv::dnn::Layer>(new ResizeBilinearLayer(params));
|
||||
}
|
||||
|
||||
virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
|
||||
const int,
|
||||
std::vector<std::vector<int> > &outputs,
|
||||
std::vector<std::vector<int> > &) const CV_OVERRIDE
|
||||
{
|
||||
std::vector<int> outShape(4);
|
||||
outShape[0] = inputs[0][0]; // batch size
|
||||
outShape[1] = inputs[0][1]; // number of channels
|
||||
outShape[2] = outHeight != 0 ? outHeight : (inputs[0][2] * factorHeight);
|
||||
outShape[3] = outWidth != 0 ? outWidth : (inputs[0][3] * factorWidth);
|
||||
outputs.assign(1, outShape);
|
||||
return false;
|
||||
}
|
||||
|
||||
virtual void finalize(cv::InputArrayOfArrays, cv::OutputArrayOfArrays outputs_arr) CV_OVERRIDE
|
||||
{
|
||||
std::vector<cv::Mat> outputs;
|
||||
outputs_arr.getMatVector(outputs);
|
||||
if (!outWidth && !outHeight)
|
||||
{
|
||||
outHeight = outputs[0].size[2];
|
||||
outWidth = outputs[0].size[3];
|
||||
}
|
||||
}
|
||||
|
||||
// This implementation is based on a reference implementation from
|
||||
// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
|
||||
virtual void forward(cv::InputArrayOfArrays inputs_arr,
|
||||
cv::OutputArrayOfArrays outputs_arr,
|
||||
cv::OutputArrayOfArrays internals_arr) CV_OVERRIDE
|
||||
{
|
||||
if (inputs_arr.depth() == CV_16S)
|
||||
{
|
||||
// In case of DNN_TARGET_OPENCL_FP16 target the following method
|
||||
// converts data from FP16 to FP32 and calls this forward again.
|
||||
forward_fallback(inputs_arr, outputs_arr, internals_arr);
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<cv::Mat> inputs, outputs;
|
||||
inputs_arr.getMatVector(inputs);
|
||||
outputs_arr.getMatVector(outputs);
|
||||
|
||||
cv::Mat& inp = inputs[0];
|
||||
cv::Mat& out = outputs[0];
|
||||
const float* inpData = (float*)inp.data;
|
||||
float* outData = (float*)out.data;
|
||||
|
||||
const int batchSize = inp.size[0];
|
||||
const int numChannels = inp.size[1];
|
||||
const int inpHeight = inp.size[2];
|
||||
const int inpWidth = inp.size[3];
|
||||
|
||||
float heightScale = static_cast<float>(inpHeight) / outHeight;
|
||||
float widthScale = static_cast<float>(inpWidth) / outWidth;
|
||||
for (int b = 0; b < batchSize; ++b)
|
||||
{
|
||||
for (int y = 0; y < outHeight; ++y)
|
||||
{
|
||||
float input_y = y * heightScale;
|
||||
int y0 = static_cast<int>(std::floor(input_y));
|
||||
int y1 = std::min(y0 + 1, inpHeight - 1);
|
||||
for (int x = 0; x < outWidth; ++x)
|
||||
{
|
||||
float input_x = x * widthScale;
|
||||
int x0 = static_cast<int>(std::floor(input_x));
|
||||
int x1 = std::min(x0 + 1, inpWidth - 1);
|
||||
for (int c = 0; c < numChannels; ++c)
|
||||
{
|
||||
float interpolation =
|
||||
inpData[offset(inp.size, c, x0, y0, b)] * (1 - (input_y - y0)) * (1 - (input_x - x0)) +
|
||||
inpData[offset(inp.size, c, x0, y1, b)] * (input_y - y0) * (1 - (input_x - x0)) +
|
||||
inpData[offset(inp.size, c, x1, y0, b)] * (1 - (input_y - y0)) * (input_x - x0) +
|
||||
inpData[offset(inp.size, c, x1, y1, b)] * (input_y - y0) * (input_x - x0);
|
||||
outData[offset(out.size, c, x, y, b)] = interpolation;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
static inline int offset(const cv::MatSize& size, int c, int x, int y, int b)
|
||||
{
|
||||
return x + size[3] * (y + size[2] * (c + size[1] * b));
|
||||
}
|
||||
|
||||
int outWidth, outHeight, factorWidth, factorHeight;
|
||||
};
|
||||
//! [ResizeBilinearLayer]
|
||||
|
||||
//
|
||||
// The following code is used only to generate tutorials documentation.
|
||||
//
|
||||
|
||||
//! [A custom layer interface]
|
||||
class MyLayer : public cv::dnn::Layer
|
||||
{
|
||||
public:
|
||||
//! [MyLayer::MyLayer]
|
||||
MyLayer(const cv::dnn::LayerParams ¶ms);
|
||||
//! [MyLayer::MyLayer]
|
||||
|
||||
//! [MyLayer::create]
|
||||
static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params);
|
||||
//! [MyLayer::create]
|
||||
|
||||
//! [MyLayer::getMemoryShapes]
|
||||
virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
|
||||
const int requiredOutputs,
|
||||
std::vector<std::vector<int> > &outputs,
|
||||
std::vector<std::vector<int> > &internals) const CV_OVERRIDE;
|
||||
//! [MyLayer::getMemoryShapes]
|
||||
|
||||
//! [MyLayer::forward]
|
||||
virtual void forward(cv::InputArrayOfArrays inputs,
|
||||
cv::OutputArrayOfArrays outputs,
|
||||
cv::OutputArrayOfArrays internals) CV_OVERRIDE;
|
||||
//! [MyLayer::forward]
|
||||
|
||||
//! [MyLayer::finalize]
|
||||
virtual void finalize(cv::InputArrayOfArrays inputs,
|
||||
cv::OutputArrayOfArrays outputs) CV_OVERRIDE;
|
||||
//! [MyLayer::finalize]
|
||||
};
|
||||
//! [A custom layer interface]
|
||||
|
||||
//! [Register a custom layer]
|
||||
#include <opencv2/dnn/layer.details.hpp> // CV_DNN_REGISTER_LAYER_CLASS
|
||||
|
||||
static inline void loadNet()
|
||||
{
|
||||
CV_DNN_REGISTER_LAYER_CLASS(Interp, InterpLayer);
|
||||
// ...
|
||||
//! [Register a custom layer]
|
||||
|
||||
//! [Register InterpLayer]
|
||||
CV_DNN_REGISTER_LAYER_CLASS(Interp, InterpLayer);
|
||||
cv::dnn::Net caffeNet = cv::dnn::readNet("/path/to/config.prototxt", "/path/to/weights.caffemodel");
|
||||
//! [Register InterpLayer]
|
||||
|
||||
//! [Register ResizeBilinearLayer]
|
||||
CV_DNN_REGISTER_LAYER_CLASS(ResizeBilinear, ResizeBilinearLayer);
|
||||
cv::dnn::Net tfNet = cv::dnn::readNet("/path/to/graph.pb");
|
||||
//! [Register ResizeBilinearLayer]
|
||||
|
||||
if (false) loadNet(); // To prevent unused function warning.
|
||||
}
|
||||
|
||||
#endif // __OPENCV_SAMPLES_DNN_CUSTOM_LAYERS__
|
||||
@@ -0,0 +1,136 @@
|
||||
/*
|
||||
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 deblurs the given blurry image.
|
||||
|
||||
Copyright (C) 2025, Bigvision LLC.
|
||||
|
||||
How to use:
|
||||
Sample command to run:
|
||||
`./example_dnn_deblurring`
|
||||
|
||||
You can download NAFNet deblurring model using
|
||||
`python download_models.py NAFNet`
|
||||
|
||||
References:
|
||||
Github: https://github.com/megvii-research/NAFNet
|
||||
PyTorch model: https://drive.google.com/file/d/14D4V4raNYIOhETfcuuLI3bGLB-OYIv6X/view
|
||||
|
||||
PyTorch model was converted to ONNX and then ONNX model was further quantized using block quantization from [opencv_zoo](https://github.com/opencv/opencv_zoo/blob/main/tools/quantize/block_quantize.py)
|
||||
|
||||
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 deblurring using OpenCV. \n\n"
|
||||
"Firstly, download required models i.e. NAFNet 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_deblurring [--input=<image_name>] \n\n"
|
||||
"Deblurring model path can also be specified using --model argument.\n\n";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | show help message}"
|
||||
"{ @alias | NAFNet | 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 | licenseplate_motion.jpg | 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;
|
||||
|
||||
|
||||
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 deblurring 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");
|
||||
float scale = parser.get<float>("scale");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
Scalar mean_v = parser.get<Scalar>("mean");
|
||||
|
||||
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));
|
||||
|
||||
Mat inputImage = imread(findFile(imgPath));
|
||||
if (inputImage.empty()) {
|
||||
cerr << "Error: Input image could not be loaded." << endl;
|
||||
return -1;
|
||||
}
|
||||
Mat image = inputImage.clone();
|
||||
|
||||
Mat image_blob = blobFromImage(image, scale, Size(image.cols, image.rows), mean_v, swapRB, false);
|
||||
|
||||
net.setInput(image_blob);
|
||||
Mat output = net.forward();
|
||||
|
||||
// Post Processing
|
||||
Mat output_transposed(3, &output.size[1], CV_32F, output.ptr<float>());
|
||||
|
||||
vector<Mat> channels = {
|
||||
Mat(output_transposed.size[1], output_transposed.size[2], CV_32F, output_transposed.ptr<float>(2)),
|
||||
Mat(output_transposed.size[1], output_transposed.size[2], CV_32F, output_transposed.ptr<float>(1)),
|
||||
Mat(output_transposed.size[1], output_transposed.size[2], CV_32F, output_transposed.ptr<float>(0))
|
||||
};
|
||||
|
||||
Mat outputImage;
|
||||
merge(channels, outputImage);
|
||||
outputImage.convertTo(outputImage, CV_8UC3, 255.0);
|
||||
|
||||
imshow("Input Image", inputImage);
|
||||
imshow("Output Image", outputImage);
|
||||
waitKey(0);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
#!/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 deblurs the given blurry image.
|
||||
|
||||
Copyright (C) 2025, Bigvision LLC.
|
||||
|
||||
How to use:
|
||||
Sample command to run:
|
||||
`python deblurring.py`
|
||||
|
||||
You can download NAFNet deblurring model using
|
||||
`python download_models.py NAFNet`
|
||||
|
||||
References:
|
||||
Github: https://github.com/megvii-research/NAFNet
|
||||
PyTorch model: https://drive.google.com/file/d/14D4V4raNYIOhETfcuuLI3bGLB-OYIv6X/view
|
||||
|
||||
PyTorch model was converted to ONNX and then ONNX model was further quantized using block quantization from [opencv_zoo](https://github.com/opencv/opencv_zoo/blob/main/tools/quantize/block_quantize.py)
|
||||
|
||||
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 cv2 as cv
|
||||
import numpy as np
|
||||
from common import *
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Use this script for image deblurring using OpenCV.
|
||||
|
||||
Firstly, download required models i.e. NAFNet 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 deblurring.py [--input=<image_name>]
|
||||
|
||||
Deblurring model path can also be specified using --model argument.
|
||||
'''
|
||||
)
|
||||
|
||||
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="licenseplate_motion.jpg", 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, 'deblurring', prefix="", alias="NAFNet")
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Image deblurring using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args()
|
||||
|
||||
def main():
|
||||
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))
|
||||
image = input_image.copy()
|
||||
height, width = image.shape[:2]
|
||||
|
||||
image_blob = cv.dnn.blobFromImage(image, args.scale, (width, height), args.mean, args.rgb, False)
|
||||
net.setInput(image_blob)
|
||||
out = net.forward()
|
||||
|
||||
# Postprocessing
|
||||
output = out[0]
|
||||
output = np.transpose(output, (1, 2, 0))
|
||||
output = np.clip(output * 255.0, 0, 255).astype(np.uint8)
|
||||
out_image = cv.cvtColor(output, cv.COLOR_RGB2BGR)
|
||||
|
||||
cv.imshow("input image: ", input_image)
|
||||
cv.imshow("output image: ", out_image)
|
||||
cv.waitKey(0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args_parser()
|
||||
main()
|
||||
@@ -0,0 +1,23 @@
|
||||
from abc import ABC, ABCMeta, abstractmethod
|
||||
|
||||
|
||||
class AbstractModel(ABC):
|
||||
|
||||
@abstractmethod
|
||||
def get_prepared_models(self):
|
||||
pass
|
||||
|
||||
|
||||
class Framework(object):
|
||||
in_blob_name = ''
|
||||
out_blob_name = ''
|
||||
|
||||
__metaclass__ = ABCMeta
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_output(self, input_blob):
|
||||
pass
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ...utils import get_final_summary_info
|
||||
|
||||
|
||||
class ClsAccEvaluation:
|
||||
log = sys.stdout
|
||||
img_classes = {}
|
||||
batch_size = 0
|
||||
|
||||
def __init__(self, log_path, img_classes_file, batch_size):
|
||||
self.log = open(log_path, 'w')
|
||||
self.img_classes = self.read_classes(img_classes_file)
|
||||
self.batch_size = batch_size
|
||||
|
||||
# collect the accuracies for both models
|
||||
self.general_quality_metric = []
|
||||
self.general_inference_time = []
|
||||
|
||||
@staticmethod
|
||||
def read_classes(img_classes_file):
|
||||
result = {}
|
||||
with open(img_classes_file) as file:
|
||||
for l in file.readlines():
|
||||
result[l.split()[0]] = int(l.split()[1])
|
||||
return result
|
||||
|
||||
def get_correct_answers(self, img_list, net_output_blob):
|
||||
correct_answers = 0
|
||||
for i in range(len(img_list)):
|
||||
indexes = np.argsort(net_output_blob[i])[-5:]
|
||||
correct_index = self.img_classes[img_list[i]]
|
||||
if correct_index in indexes:
|
||||
correct_answers += 1
|
||||
return correct_answers
|
||||
|
||||
def process(self, frameworks, data_fetcher):
|
||||
sorted_imgs_names = sorted(self.img_classes.keys())
|
||||
correct_answers = [0] * len(frameworks)
|
||||
samples_handled = 0
|
||||
blobs_l1_diff = [0] * len(frameworks)
|
||||
blobs_l1_diff_count = [0] * len(frameworks)
|
||||
blobs_l_inf_diff = [sys.float_info.min] * len(frameworks)
|
||||
inference_time = [0.0] * len(frameworks)
|
||||
|
||||
for x in range(0, len(sorted_imgs_names), self.batch_size):
|
||||
sublist = sorted_imgs_names[x:x + self.batch_size]
|
||||
batch = data_fetcher.get_batch(sublist)
|
||||
|
||||
samples_handled += len(sublist)
|
||||
fw_accuracy = []
|
||||
fw_time = []
|
||||
frameworks_out = []
|
||||
for i in range(len(frameworks)):
|
||||
start = time.time()
|
||||
out = frameworks[i].get_output(batch)
|
||||
end = time.time()
|
||||
correct_answers[i] += self.get_correct_answers(sublist, out)
|
||||
fw_accuracy.append(100 * correct_answers[i] / float(samples_handled))
|
||||
frameworks_out.append(out)
|
||||
inference_time[i] += end - start
|
||||
fw_time.append(inference_time[i] / samples_handled * 1000)
|
||||
print(samples_handled, 'Accuracy for', frameworks[i].get_name() + ':', fw_accuracy[i], file=self.log)
|
||||
print("Inference time, ms ", frameworks[i].get_name(), fw_time[i], file=self.log)
|
||||
|
||||
self.general_quality_metric.append(fw_accuracy)
|
||||
self.general_inference_time.append(fw_time)
|
||||
|
||||
for i in range(1, len(frameworks)):
|
||||
log_str = frameworks[0].get_name() + " vs " + frameworks[i].get_name() + ':'
|
||||
diff = np.abs(frameworks_out[0] - frameworks_out[i])
|
||||
l1_diff = np.sum(diff) / diff.size
|
||||
print(samples_handled, "L1 difference", log_str, l1_diff, file=self.log)
|
||||
blobs_l1_diff[i] += l1_diff
|
||||
blobs_l1_diff_count[i] += 1
|
||||
if np.max(diff) > blobs_l_inf_diff[i]:
|
||||
blobs_l_inf_diff[i] = np.max(diff)
|
||||
print(samples_handled, "L_INF difference", log_str, blobs_l_inf_diff[i], file=self.log)
|
||||
|
||||
self.log.flush()
|
||||
|
||||
for i in range(1, len(blobs_l1_diff)):
|
||||
log_str = frameworks[0].get_name() + " vs " + frameworks[i].get_name() + ':'
|
||||
print('Final l1 diff', log_str, blobs_l1_diff[i] / blobs_l1_diff_count[i], file=self.log)
|
||||
|
||||
print(
|
||||
get_final_summary_info(
|
||||
self.general_quality_metric,
|
||||
self.general_inference_time,
|
||||
"accuracy"
|
||||
),
|
||||
file=self.log
|
||||
)
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
import os
|
||||
from abc import ABCMeta, abstractmethod
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ...img_utils import read_rgb_img, get_pytorch_preprocess
|
||||
from ...test.configs.default_preprocess_config import PYTORCH_RSZ_HEIGHT, PYTORCH_RSZ_WIDTH
|
||||
|
||||
|
||||
class DataFetch(object):
|
||||
imgs_dir = ''
|
||||
frame_size = 0
|
||||
bgr_to_rgb = False
|
||||
|
||||
__metaclass__ = ABCMeta
|
||||
|
||||
@abstractmethod
|
||||
def preprocess(self, img):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def reshape_img(img):
|
||||
img = img[:, :, 0:3].transpose(2, 0, 1)
|
||||
return np.expand_dims(img, 0)
|
||||
|
||||
def center_crop(self, img):
|
||||
cols = img.shape[1]
|
||||
rows = img.shape[0]
|
||||
|
||||
y1 = round((rows - self.frame_size) / 2)
|
||||
y2 = round(y1 + self.frame_size)
|
||||
x1 = round((cols - self.frame_size) / 2)
|
||||
x2 = round(x1 + self.frame_size)
|
||||
return img[y1:y2, x1:x2]
|
||||
|
||||
def initial_preprocess(self, img):
|
||||
min_dim = min(img.shape[-3], img.shape[-2])
|
||||
resize_ratio = self.frame_size / float(min_dim)
|
||||
|
||||
img = cv2.resize(img, (0, 0), fx=resize_ratio, fy=resize_ratio)
|
||||
img = self.center_crop(img)
|
||||
return img
|
||||
|
||||
def get_preprocessed_img(self, img_path):
|
||||
image_data = read_rgb_img(img_path, self.bgr_to_rgb)
|
||||
image_data = self.preprocess(image_data)
|
||||
return self.reshape_img(image_data)
|
||||
|
||||
def get_batch(self, img_names):
|
||||
assert type(img_names) is list
|
||||
batch = np.zeros((len(img_names), 3, self.frame_size, self.frame_size)).astype(np.float32)
|
||||
|
||||
for i in range(len(img_names)):
|
||||
img_name = img_names[i]
|
||||
img_file = os.path.join(self.imgs_dir, img_name)
|
||||
assert os.path.exists(img_file)
|
||||
|
||||
batch[i] = self.get_preprocessed_img(img_file)
|
||||
return batch
|
||||
|
||||
|
||||
class PyTorchPreprocessedFetch(DataFetch):
|
||||
def __init__(self, pytorch_cls_config, preprocess_input=None):
|
||||
self.imgs_dir = pytorch_cls_config.img_root_dir
|
||||
self.frame_size = pytorch_cls_config.frame_size
|
||||
self.bgr_to_rgb = pytorch_cls_config.bgr_to_rgb
|
||||
self.preprocess_input = preprocess_input
|
||||
|
||||
def preprocess(self, img):
|
||||
img = cv2.resize(img, (PYTORCH_RSZ_WIDTH, PYTORCH_RSZ_HEIGHT))
|
||||
img = self.center_crop(img)
|
||||
if self.preprocess_input:
|
||||
return self.presprocess_input(img)
|
||||
return get_pytorch_preprocess(img)
|
||||
|
||||
|
||||
class TFPreprocessedFetch(DataFetch):
|
||||
def __init__(self, tf_cls_config, preprocess_input):
|
||||
self.imgs_dir = tf_cls_config.img_root_dir
|
||||
self.frame_size = tf_cls_config.frame_size
|
||||
self.bgr_to_rgb = tf_cls_config.bgr_to_rgb
|
||||
self.preprocess_input = preprocess_input
|
||||
|
||||
def preprocess(self, img):
|
||||
img = self.initial_preprocess(img)
|
||||
return self.preprocess_input(img)
|
||||
@@ -0,0 +1,19 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .test.configs.default_preprocess_config import BASE_IMG_SCALE_FACTOR
|
||||
|
||||
|
||||
def read_rgb_img(img_file, is_bgr_to_rgb=True):
|
||||
img = cv2.imread(img_file, cv2.IMREAD_COLOR)
|
||||
if is_bgr_to_rgb:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
return img
|
||||
|
||||
|
||||
def get_pytorch_preprocess(img):
|
||||
img = img.astype(np.float32)
|
||||
img *= BASE_IMG_SCALE_FACTOR
|
||||
img -= [0.485, 0.456, 0.406]
|
||||
img /= [0.229, 0.224, 0.225]
|
||||
return img
|
||||
@@ -0,0 +1,60 @@
|
||||
from .configs.test_config import TestClsConfig, TestClsModuleConfig
|
||||
from .model_test_pipeline import ModelTestPipeline
|
||||
from ..evaluation.classification.cls_accuracy_evaluator import ClsAccEvaluation
|
||||
from ..utils import get_test_module
|
||||
|
||||
|
||||
class ClsModelTestPipeline(ModelTestPipeline):
|
||||
def __init__(
|
||||
self,
|
||||
network_model,
|
||||
model_processor,
|
||||
dnn_model_processor,
|
||||
data_fetcher,
|
||||
img_processor=None,
|
||||
cls_args_parser=None,
|
||||
default_input_blob_preproc=None
|
||||
):
|
||||
super(ClsModelTestPipeline, self).__init__(
|
||||
network_model,
|
||||
model_processor,
|
||||
dnn_model_processor
|
||||
)
|
||||
|
||||
if cls_args_parser:
|
||||
self._parser = cls_args_parser
|
||||
|
||||
self.test_config = TestClsConfig()
|
||||
|
||||
parser_args = self._parser.parse_args()
|
||||
|
||||
if parser_args.test:
|
||||
self._test_module_config = TestClsModuleConfig()
|
||||
self._test_module = get_test_module(
|
||||
self._test_module_config.test_module_name,
|
||||
self._test_module_config.test_module_path
|
||||
)
|
||||
|
||||
if parser_args.default_img_preprocess:
|
||||
self._default_input_blob_preproc = default_input_blob_preproc
|
||||
if parser_args.evaluate:
|
||||
self._data_fetcher = data_fetcher(self.test_config, img_processor)
|
||||
|
||||
def _configure_test_module_params(self):
|
||||
self._test_module_param_list.extend((
|
||||
'--crop', self._test_module_config.crop,
|
||||
'--std', *self._test_module_config.std
|
||||
))
|
||||
|
||||
if self._test_module_config.rsz_height and self._test_module_config.rsz_width:
|
||||
self._test_module_param_list.extend((
|
||||
'--initial_height', self._test_module_config.rsz_height,
|
||||
'--initial_width', self._test_module_config.rsz_width,
|
||||
))
|
||||
|
||||
def _configure_acc_eval(self, log_path):
|
||||
self._accuracy_evaluator = ClsAccEvaluation(
|
||||
log_path,
|
||||
self.test_config.img_cls_file,
|
||||
self.test_config.batch_size
|
||||
)
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
BASE_IMG_SCALE_FACTOR = 1 / 255.0
|
||||
PYTORCH_RSZ_HEIGHT = 256
|
||||
PYTORCH_RSZ_WIDTH = 256
|
||||
|
||||
pytorch_resize_input_blob = {
|
||||
"mean": ["123.675", "116.28", "103.53"],
|
||||
"scale": str(BASE_IMG_SCALE_FACTOR),
|
||||
"std": ["0.229", "0.224", "0.225"],
|
||||
"crop": "True",
|
||||
"rgb": True,
|
||||
"rsz_height": str(PYTORCH_RSZ_HEIGHT),
|
||||
"rsz_width": str(PYTORCH_RSZ_WIDTH)
|
||||
}
|
||||
|
||||
pytorch_input_blob = {
|
||||
"mean": ["123.675", "116.28", "103.53"],
|
||||
"scale": str(BASE_IMG_SCALE_FACTOR),
|
||||
"std": ["0.229", "0.224", "0.225"],
|
||||
"crop": "True",
|
||||
"rgb": True
|
||||
}
|
||||
|
||||
tf_input_blob = {
|
||||
"scale": str(1 / 127.5),
|
||||
"mean": ["127.5", "127.5", "127.5"],
|
||||
"std": [],
|
||||
"crop": "True",
|
||||
"rgb": True
|
||||
}
|
||||
|
||||
tf_model_blob_caffe_mode = {
|
||||
"mean": ["103.939", "116.779", "123.68"],
|
||||
"scale": "1.0",
|
||||
"std": [],
|
||||
"crop": "True",
|
||||
"rgb": False
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommonConfig:
|
||||
output_data_root_dir: str = "dnn_model_runner/dnn_conversion"
|
||||
logs_dir: str = os.path.join(output_data_root_dir, "logs")
|
||||
log_file_path: str = os.path.join(logs_dir, "{}_log.txt")
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestClsConfig:
|
||||
batch_size: int = 1
|
||||
frame_size: int = 224
|
||||
img_root_dir: str = "./ILSVRC2012_img_val"
|
||||
# location of image-class matching
|
||||
img_cls_file: str = "./val.txt"
|
||||
bgr_to_rgb: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestClsModuleConfig:
|
||||
cls_test_data_dir: str = "../data"
|
||||
test_module_name: str = "classification"
|
||||
test_module_path: str = "classification.py"
|
||||
input_img: str = os.path.join(cls_test_data_dir, "squirrel_cls.jpg")
|
||||
model: str = ""
|
||||
|
||||
frame_height: str = str(TestClsConfig.frame_size)
|
||||
frame_width: str = str(TestClsConfig.frame_size)
|
||||
scale: str = "1.0"
|
||||
mean: List[str] = field(default_factory=lambda: ["0.0", "0.0", "0.0"])
|
||||
std: List[str] = field(default_factory=list)
|
||||
crop: str = "False"
|
||||
rgb: str = "True"
|
||||
rsz_height: str = ""
|
||||
rsz_width: str = ""
|
||||
classes: str = os.path.join(cls_test_data_dir, "dnn", "classification_classes_ILSVRC2012.txt")
|
||||
@@ -0,0 +1,126 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .configs.test_config import CommonConfig
|
||||
from ..utils import create_parser, plot_acc
|
||||
|
||||
|
||||
class ModelTestPipeline:
|
||||
def __init__(
|
||||
self,
|
||||
network_model,
|
||||
model_processor,
|
||||
dnn_model_processor
|
||||
):
|
||||
self._net_model = network_model
|
||||
self._model_processor = model_processor
|
||||
self._dnn_model_processor = dnn_model_processor
|
||||
|
||||
self._parser = create_parser()
|
||||
|
||||
self._test_module = None
|
||||
self._test_module_config = None
|
||||
self._test_module_param_list = None
|
||||
|
||||
self.test_config = None
|
||||
self._data_fetcher = None
|
||||
|
||||
self._default_input_blob_preproc = None
|
||||
self._accuracy_evaluator = None
|
||||
|
||||
def init_test_pipeline(self):
|
||||
cmd_args = self._parser.parse_args()
|
||||
model_dict = self._net_model.get_prepared_models()
|
||||
|
||||
model_names = list(model_dict.keys())
|
||||
print(
|
||||
"The model {} was successfully obtained and converted to OpenCV {}".format(model_names[0], model_names[1])
|
||||
)
|
||||
|
||||
if cmd_args.test:
|
||||
if not self._test_module_config.model:
|
||||
self._test_module_config.model = self._net_model.model_path["full_path"]
|
||||
|
||||
if cmd_args.default_img_preprocess:
|
||||
self._test_module_config.scale = self._default_input_blob_preproc["scale"]
|
||||
self._test_module_config.mean = self._default_input_blob_preproc["mean"]
|
||||
self._test_module_config.std = self._default_input_blob_preproc["std"]
|
||||
self._test_module_config.crop = self._default_input_blob_preproc["crop"]
|
||||
|
||||
if "rsz_height" in self._default_input_blob_preproc and "rsz_width" in self._default_input_blob_preproc:
|
||||
self._test_module_config.rsz_height = self._default_input_blob_preproc["rsz_height"]
|
||||
self._test_module_config.rsz_width = self._default_input_blob_preproc["rsz_width"]
|
||||
|
||||
self._test_module_param_list = [
|
||||
'--model', self._test_module_config.model,
|
||||
'--input', self._test_module_config.input_img,
|
||||
'--width', self._test_module_config.frame_width,
|
||||
'--height', self._test_module_config.frame_height,
|
||||
'--scale', self._test_module_config.scale,
|
||||
'--mean', *self._test_module_config.mean,
|
||||
'--std', *self._test_module_config.std,
|
||||
'--classes', self._test_module_config.classes,
|
||||
]
|
||||
|
||||
if self._default_input_blob_preproc["rgb"]:
|
||||
self._test_module_param_list.append('--rgb')
|
||||
|
||||
self._configure_test_module_params()
|
||||
|
||||
self._test_module.main(
|
||||
self._test_module_param_list
|
||||
)
|
||||
|
||||
if cmd_args.evaluate:
|
||||
original_model_name = model_names[0]
|
||||
dnn_model_name = model_names[1]
|
||||
|
||||
self.run_test_pipeline(
|
||||
[
|
||||
self._model_processor(model_dict[original_model_name], original_model_name),
|
||||
self._dnn_model_processor(model_dict[dnn_model_name], dnn_model_name)
|
||||
],
|
||||
original_model_name.replace(" ", "_")
|
||||
)
|
||||
|
||||
def run_test_pipeline(
|
||||
self,
|
||||
models_list,
|
||||
formatted_exp_name,
|
||||
is_plot_acc=True
|
||||
):
|
||||
log_path, logs_dir = self._configure_eval_log(formatted_exp_name)
|
||||
|
||||
print(
|
||||
"===== Running evaluation of the model with the following params:\n"
|
||||
"\t* val data location: {}\n"
|
||||
"\t* log file location: {}\n".format(
|
||||
self.test_config.img_root_dir,
|
||||
log_path
|
||||
)
|
||||
)
|
||||
|
||||
os.makedirs(logs_dir, exist_ok=True)
|
||||
|
||||
self._configure_acc_eval(log_path)
|
||||
self._accuracy_evaluator.process(models_list, self._data_fetcher)
|
||||
|
||||
if is_plot_acc:
|
||||
plot_acc(
|
||||
np.array(self._accuracy_evaluator.general_inference_time),
|
||||
formatted_exp_name
|
||||
)
|
||||
|
||||
print("===== End of the evaluation pipeline =====")
|
||||
|
||||
def _configure_acc_eval(self, log_path):
|
||||
pass
|
||||
|
||||
def _configure_test_module_params(self):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _configure_eval_log(formatted_exp_name):
|
||||
common_test_config = CommonConfig()
|
||||
return common_test_config.log_file_path.format(formatted_exp_name), common_test_config.logs_dir
|
||||
@@ -0,0 +1,153 @@
|
||||
import argparse
|
||||
import importlib.util
|
||||
import os
|
||||
import random
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
import torch
|
||||
|
||||
from .test.configs.test_config import CommonConfig
|
||||
|
||||
SEED_VAL = 42
|
||||
DNN_LIB = "DNN"
|
||||
# common path for model savings
|
||||
MODEL_PATH_ROOT = os.path.join(CommonConfig().output_data_root_dir, "{}/models")
|
||||
|
||||
|
||||
def get_full_model_path(lib_name, model_full_name):
|
||||
model_path = MODEL_PATH_ROOT.format(lib_name)
|
||||
return {
|
||||
"path": model_path,
|
||||
"full_path": os.path.join(model_path, model_full_name)
|
||||
}
|
||||
|
||||
|
||||
def plot_acc(data_list, experiment_name):
|
||||
plt.figure(figsize=[8, 6])
|
||||
plt.plot(data_list[:, 0], "r", linewidth=2.5, label="Original Model")
|
||||
plt.plot(data_list[:, 1], "b", linewidth=2.5, label="Converted DNN Model")
|
||||
plt.xlabel("Iterations ", fontsize=15)
|
||||
plt.ylabel("Time (ms)", fontsize=15)
|
||||
plt.title(experiment_name, fontsize=15)
|
||||
plt.legend()
|
||||
full_path_to_fig = os.path.join(CommonConfig().output_data_root_dir, experiment_name + ".png")
|
||||
plt.savefig(full_path_to_fig, bbox_inches="tight")
|
||||
|
||||
|
||||
def get_final_summary_info(general_quality_metric, general_inference_time, metric_name):
|
||||
general_quality_metric = np.array(general_quality_metric)
|
||||
general_inference_time = np.array(general_inference_time)
|
||||
summary_line = "===== End of processing. General results:\n"
|
||||
"\t* mean {} for the original model: {}\t"
|
||||
"\t* mean time (min) for the original model inferences: {}\n"
|
||||
"\t* mean {} for the DNN model: {}\t"
|
||||
"\t* mean time (min) for the DNN model inferences: {}\n".format(
|
||||
metric_name, np.mean(general_quality_metric[:, 0]),
|
||||
np.mean(general_inference_time[:, 0]) / 60000,
|
||||
metric_name, np.mean(general_quality_metric[:, 1]),
|
||||
np.mean(general_inference_time[:, 1]) / 60000,
|
||||
)
|
||||
return summary_line
|
||||
|
||||
|
||||
def set_common_reproducibility():
|
||||
random.seed(SEED_VAL)
|
||||
np.random.seed(SEED_VAL)
|
||||
|
||||
|
||||
def set_pytorch_env():
|
||||
set_common_reproducibility()
|
||||
torch.manual_seed(SEED_VAL)
|
||||
torch.set_printoptions(precision=10)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(SEED_VAL)
|
||||
torch.backends.cudnn_benchmark_enabled = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
|
||||
def set_tf_env(is_use_gpu=True):
|
||||
set_common_reproducibility()
|
||||
tf.random.set_seed(SEED_VAL)
|
||||
os.environ["TF_DETERMINISTIC_OPS"] = "1"
|
||||
|
||||
if tf.config.list_physical_devices("GPU") and is_use_gpu:
|
||||
gpu_devices = tf.config.list_physical_devices("GPU")
|
||||
tf.config.experimental.set_visible_devices(gpu_devices[0], "GPU")
|
||||
tf.config.experimental.set_memory_growth(gpu_devices[0], True)
|
||||
os.environ["TF_USE_CUDNN"] = "1"
|
||||
else:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
||||
|
||||
|
||||
def str_bool(input_val):
|
||||
if input_val.lower() in ('yes', 'true', 't', 'y', '1'):
|
||||
return True
|
||||
elif input_val.lower() in ('no', 'false', 'f', 'n', '0'):
|
||||
return False
|
||||
else:
|
||||
raise argparse.ArgumentTypeError('Boolean value was expected')
|
||||
|
||||
|
||||
def get_formatted_model_list(model_list):
|
||||
note_line = 'Please, choose the model from the below list:\n'
|
||||
spaces_to_set = ' ' * (len(note_line) - 2)
|
||||
return note_line + ''.join([spaces_to_set, '{} \n'] * len(model_list)).format(*model_list)
|
||||
|
||||
|
||||
def model_str(model_list):
|
||||
def type_model_list(input_val):
|
||||
if input_val.lower() in model_list:
|
||||
return input_val.lower()
|
||||
else:
|
||||
raise argparse.ArgumentTypeError(
|
||||
'The model is currently unavailable for test.\n' +
|
||||
get_formatted_model_list(model_list)
|
||||
)
|
||||
|
||||
return type_model_list
|
||||
|
||||
|
||||
def get_test_module(test_module_name, test_module_path):
|
||||
module_spec = importlib.util.spec_from_file_location(test_module_name, test_module_path)
|
||||
test_module = importlib.util.module_from_spec(module_spec)
|
||||
module_spec.loader.exec_module(test_module)
|
||||
module_spec.loader.exec_module(test_module)
|
||||
return test_module
|
||||
|
||||
|
||||
def create_parser():
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
|
||||
parser.add_argument(
|
||||
"--test",
|
||||
type=str_bool,
|
||||
help="Define whether you'd like to run the model with OpenCV for testing.",
|
||||
default=False
|
||||
),
|
||||
parser.add_argument(
|
||||
"--default_img_preprocess",
|
||||
type=str_bool,
|
||||
help="Define whether you'd like to preprocess the input image with defined"
|
||||
" PyTorch or TF functions for model test with OpenCV.",
|
||||
default=False
|
||||
),
|
||||
parser.add_argument(
|
||||
"--evaluate",
|
||||
type=str_bool,
|
||||
help="Define whether you'd like to run evaluation of the models (ex.: TF vs OpenCV networks).",
|
||||
default=True
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def create_extended_parser(model_list):
|
||||
parser = create_parser()
|
||||
parser.add_argument(
|
||||
"--model_name",
|
||||
type=model_str(model_list=model_list),
|
||||
help="\nDefine the model name to test.\n" +
|
||||
get_formatted_model_list(model_list),
|
||||
required=True
|
||||
)
|
||||
return parser
|
||||
@@ -0,0 +1,78 @@
|
||||
# Run PaddlePaddle model using OpenCV
|
||||
|
||||
These two demonstrations show how to inference PaddlePaddle model using OpenCV.
|
||||
|
||||
## Environment Setup
|
||||
|
||||
```shell
|
||||
pip install paddlepaddle-gpu
|
||||
pip install paddlehub
|
||||
pip install paddle2onnx
|
||||
```
|
||||
|
||||
## 1. Run PaddlePaddle ResNet50 using OpenCV
|
||||
|
||||
### Run PaddlePaddle model demo
|
||||
|
||||
Run the code sample as follows:
|
||||
|
||||
```shell
|
||||
python paddle_resnet50.py
|
||||
```
|
||||
|
||||
There are three parts to the process:
|
||||
|
||||
1. Export PaddlePaddle ResNet50 model to onnx format.
|
||||
2. Use `cv2.dnn.readNetFromONNX` to load the model file.
|
||||
3. Preprocess image file and do the inference.
|
||||
|
||||
## 2. Run PaddleSeg Portrait Segmentation using OpenCV
|
||||
|
||||
### Convert to ONNX Model
|
||||
|
||||
#### 1. Get Paddle Inference model
|
||||
|
||||
For more details, please refer to [PaddleSeg](https://github.com/PaddlePaddle/PaddleSeg/blob/release/2.1/contrib/HumanSeg/README.md).
|
||||
|
||||
```shell
|
||||
wget https://x2paddle.bj.bcebos.com/inference/models/humanseg_hrnet18_small_v1.zip
|
||||
unzip humanseg_hrnet18_small_v1.zip
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
* The exported model must have a fixed input shape, as dynamic is not supported at this moment.
|
||||
|
||||
#### 2. Convert to ONNX model using paddle2onnx
|
||||
|
||||
To convert the model, use the following command:
|
||||
|
||||
```
|
||||
paddle2onnx --model_dir humanseg_hrnet18_small_v1 \
|
||||
--model_filename model.pdmodel \
|
||||
--params_filename model.pdiparams \
|
||||
--opset_version 11 \
|
||||
--save_file humanseg_hrnet18_tiny.onnx
|
||||
```
|
||||
|
||||
The converted model can be found in the current directory by the name `humanseg_hrnet18_tiny.onnx` .
|
||||
|
||||
### Run PaddleSeg Portrait Segmentation demo
|
||||
|
||||
Run the code sample as follows:
|
||||
|
||||
```shell
|
||||
python paddle_humanseg.py
|
||||
```
|
||||
|
||||
There are three parts to the process:
|
||||
|
||||
1. Use `cv2.dnn.readNetFromONNX` to load the model file.
|
||||
2. Preprocess image file and do inference.
|
||||
3. Postprocess image file and visualize.
|
||||
|
||||
The resulting file can be found at `data/result_test_human.jpg` .
|
||||
|
||||
### Portrait segmentation visualization
|
||||
|
||||
<img src="../../../../data/messi5.jpg" width="50%" height="50%"><img src="./data/result_test_human.jpg" width="50%" height="50%">
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 126 KiB |
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 61 KiB |
@@ -0,0 +1,112 @@
|
||||
import os
|
||||
import paddlehub.vision.transforms as T
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
|
||||
|
||||
def get_color_map_list(num_classes):
|
||||
"""
|
||||
Returns the color map for visualizing the segmentation mask,
|
||||
which can support arbitrary number of classes.
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes.
|
||||
|
||||
Returns:
|
||||
(list). The color map.
|
||||
"""
|
||||
|
||||
num_classes += 1
|
||||
color_map = num_classes * [0, 0, 0]
|
||||
for i in range(0, num_classes):
|
||||
j = 0
|
||||
lab = i
|
||||
while lab:
|
||||
color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
|
||||
color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
|
||||
color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
|
||||
j += 1
|
||||
lab >>= 3
|
||||
color_map = color_map[3:]
|
||||
return color_map
|
||||
|
||||
|
||||
def visualize(image, result, save_dir=None, weight=0.6):
|
||||
"""
|
||||
Convert predict result to color image, and save added image.
|
||||
|
||||
Args:
|
||||
image (str): The path of origin image.
|
||||
result (np.ndarray): The predict result of image.
|
||||
save_dir (str): The directory for saving visual image. Default: None.
|
||||
weight (float): The image weight of visual image, and the result weight is (1 - weight). Default: 0.6
|
||||
|
||||
Returns:
|
||||
vis_result (np.ndarray): If `save_dir` is None, return the visualized result.
|
||||
"""
|
||||
|
||||
color_map = get_color_map_list(256)
|
||||
color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
|
||||
color_map = np.array(color_map).astype("uint8")
|
||||
# Use OpenCV LUT for color mapping
|
||||
c1 = cv.LUT(result, color_map[:, 0])
|
||||
c2 = cv.LUT(result, color_map[:, 1])
|
||||
c3 = cv.LUT(result, color_map[:, 2])
|
||||
pseudo_img = np.dstack((c1, c2, c3))
|
||||
|
||||
im = cv.imread(image)
|
||||
vis_result = cv.addWeighted(im, weight, pseudo_img, 1 - weight, 0)
|
||||
|
||||
if save_dir is not None:
|
||||
if not os.path.exists(save_dir):
|
||||
os.makedirs(save_dir)
|
||||
image_name = os.path.split(image)[-1]
|
||||
out_path = os.path.join(save_dir, image_name)
|
||||
cv.imwrite(out_path, vis_result)
|
||||
else:
|
||||
return vis_result
|
||||
|
||||
|
||||
def preprocess(image_path):
|
||||
''' preprocess input image file to np.ndarray
|
||||
|
||||
Args:
|
||||
image_path(str): Path of input image file
|
||||
|
||||
Returns:
|
||||
ProcessedImage(numpy.ndarray): A numpy.ndarray
|
||||
variable which shape is (1, 3, 192, 192)
|
||||
'''
|
||||
transforms = T.Compose([
|
||||
T.Resize((192, 192)),
|
||||
T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
|
||||
],
|
||||
to_rgb=True)
|
||||
return np.expand_dims(transforms(image_path), axis=0)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img_path = "../../../../data/messi5.jpg"
|
||||
# load PPSeg Model use cv.dnn
|
||||
net = cv.dnn.readNetFromONNX('humanseg_hrnet18_tiny.onnx')
|
||||
# read and preprocess image file
|
||||
im = preprocess(img_path)
|
||||
# inference
|
||||
net.setInput(im)
|
||||
result = net.forward(['save_infer_model/scale_0.tmp_1'])
|
||||
# post process
|
||||
image = cv.imread(img_path)
|
||||
r, c, _ = image.shape
|
||||
result = np.argmax(result[0], axis=1).astype(np.uint8)
|
||||
result = cv.resize(result[0, :, :],
|
||||
dsize=(c, r),
|
||||
interpolation=cv.INTER_NEAREST)
|
||||
|
||||
print("grid_image.shape is: ", result.shape)
|
||||
folder_path = "data"
|
||||
if not os.path.exists(folder_path):
|
||||
os.makedirs(folder_path)
|
||||
file_path = os.path.join(folder_path, '%s.jpg' % "result_test_human")
|
||||
result_color = visualize(img_path, result)
|
||||
cv.imwrite(file_path, result_color)
|
||||
print('%s saved' % file_path)
|
||||
@@ -0,0 +1,61 @@
|
||||
import paddle
|
||||
import paddlehub as hub
|
||||
import paddlehub.vision.transforms as T
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
|
||||
|
||||
def preprocess(image_path):
|
||||
''' preprocess input image file to np.ndarray
|
||||
|
||||
Args:
|
||||
image_path(str): Path of input image file
|
||||
|
||||
Returns:
|
||||
ProcessedImage(numpy.ndarray): A numpy.ndarray
|
||||
variable which shape is (1, 3, 224, 224)
|
||||
'''
|
||||
transforms = T.Compose([
|
||||
T.Resize((256, 256)),
|
||||
T.CenterCrop(224),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225])],
|
||||
to_rgb=True)
|
||||
return np.expand_dims(transforms(image_path), axis=0)
|
||||
|
||||
|
||||
def export_onnx_resnet50(save_path):
|
||||
''' export PaddlePaddle model to ONNX format
|
||||
|
||||
Args:
|
||||
save_path(str): Path to save exported ONNX model
|
||||
|
||||
Returns:
|
||||
None
|
||||
'''
|
||||
model = hub.Module(name="resnet50_vd_imagenet_ssld")
|
||||
input_spec = paddle.static.InputSpec(
|
||||
[1, 3, 224, 224], "float32", "image")
|
||||
paddle.onnx.export(model, save_path,
|
||||
input_spec=[input_spec],
|
||||
opset_version=10)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
save_path = './resnet50'
|
||||
image_file = './data/cat.jpg'
|
||||
labels = open('./data/labels.txt').read().strip().split('\n')
|
||||
model = export_onnx_resnet50(save_path)
|
||||
|
||||
# load resnet50 use cv.dnn
|
||||
net = cv.dnn.readNetFromONNX(save_path + '.onnx')
|
||||
# read and preprocess image file
|
||||
im = preprocess(image_file)
|
||||
# inference
|
||||
net.setInput(im)
|
||||
result = net.forward(['save_infer_model/scale_0.tmp_0'])
|
||||
# post process
|
||||
class_id = np.argmax(result[0])
|
||||
label = labels[class_id]
|
||||
print("Image: {}".format(image_file))
|
||||
print("Predict Category: {}".format(label))
|
||||
@@ -0,0 +1,71 @@
|
||||
from torchvision import models
|
||||
|
||||
from ..pytorch_model import (
|
||||
PyTorchModelPreparer,
|
||||
PyTorchModelProcessor,
|
||||
PyTorchDnnModelProcessor
|
||||
)
|
||||
from ...common.evaluation.classification.cls_data_fetcher import PyTorchPreprocessedFetch
|
||||
from ...common.test.cls_model_test_pipeline import ClsModelTestPipeline
|
||||
from ...common.test.configs.default_preprocess_config import pytorch_resize_input_blob
|
||||
from ...common.test.configs.test_config import TestClsConfig
|
||||
from ...common.utils import set_pytorch_env, create_extended_parser
|
||||
|
||||
model_dict = {
|
||||
"alexnet": models.alexnet,
|
||||
|
||||
"vgg11": models.vgg11,
|
||||
"vgg13": models.vgg13,
|
||||
"vgg16": models.vgg16,
|
||||
"vgg19": models.vgg19,
|
||||
|
||||
"resnet18": models.resnet18,
|
||||
"resnet34": models.resnet34,
|
||||
"resnet50": models.resnet50,
|
||||
"resnet101": models.resnet101,
|
||||
"resnet152": models.resnet152,
|
||||
|
||||
"squeezenet1_0": models.squeezenet1_0,
|
||||
"squeezenet1_1": models.squeezenet1_1,
|
||||
|
||||
"resnext50_32x4d": models.resnext50_32x4d,
|
||||
"resnext101_32x8d": models.resnext101_32x8d,
|
||||
|
||||
"wide_resnet50_2": models.wide_resnet50_2,
|
||||
"wide_resnet101_2": models.wide_resnet101_2
|
||||
}
|
||||
|
||||
|
||||
class PyTorchClsModel(PyTorchModelPreparer):
|
||||
def __init__(self, height, width, model_name, original_model):
|
||||
super(PyTorchClsModel, self).__init__(height, width, model_name, original_model)
|
||||
|
||||
|
||||
def main():
|
||||
set_pytorch_env()
|
||||
|
||||
parser = create_extended_parser(list(model_dict.keys()))
|
||||
cmd_args = parser.parse_args()
|
||||
model_name = cmd_args.model_name
|
||||
|
||||
cls_model = PyTorchClsModel(
|
||||
height=TestClsConfig().frame_size,
|
||||
width=TestClsConfig().frame_size,
|
||||
model_name=model_name,
|
||||
original_model=model_dict[model_name](pretrained=True)
|
||||
)
|
||||
|
||||
pytorch_cls_pipeline = ClsModelTestPipeline(
|
||||
network_model=cls_model,
|
||||
model_processor=PyTorchModelProcessor,
|
||||
dnn_model_processor=PyTorchDnnModelProcessor,
|
||||
data_fetcher=PyTorchPreprocessedFetch,
|
||||
cls_args_parser=parser,
|
||||
default_input_blob_preproc=pytorch_resize_input_blob
|
||||
)
|
||||
|
||||
pytorch_cls_pipeline.init_test_pipeline()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.onnx
|
||||
from torch.autograd import Variable
|
||||
from torchvision import models
|
||||
|
||||
|
||||
def get_pytorch_onnx_model(original_model):
|
||||
# define the directory for further converted model save
|
||||
onnx_model_path = "models"
|
||||
# define the name of further converted model
|
||||
onnx_model_name = "resnet50.onnx"
|
||||
|
||||
# create directory for further converted model
|
||||
os.makedirs(onnx_model_path, exist_ok=True)
|
||||
|
||||
# get full path to the converted model
|
||||
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
|
||||
|
||||
# generate model input
|
||||
generated_input = Variable(
|
||||
torch.randn(1, 3, 224, 224)
|
||||
)
|
||||
|
||||
# model export into ONNX format
|
||||
torch.onnx.export(
|
||||
original_model,
|
||||
generated_input,
|
||||
full_model_path,
|
||||
verbose=True,
|
||||
input_names=["input"],
|
||||
output_names=["output"],
|
||||
opset_version=11
|
||||
)
|
||||
|
||||
return full_model_path
|
||||
|
||||
|
||||
def get_preprocessed_img(img_path):
|
||||
# read the image
|
||||
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
|
||||
input_img = input_img.astype(np.float32)
|
||||
|
||||
input_img = cv2.resize(input_img, (256, 256))
|
||||
|
||||
# define preprocess parameters
|
||||
mean = np.array([0.485, 0.456, 0.406]) * 255.0
|
||||
scale = 1 / 255.0
|
||||
std = [0.229, 0.224, 0.225]
|
||||
|
||||
# prepare input blob to fit the model input:
|
||||
# 1. subtract mean
|
||||
# 2. scale to set pixel values from 0 to 1
|
||||
input_blob = cv2.dnn.blobFromImage(
|
||||
image=input_img,
|
||||
scalefactor=scale,
|
||||
size=(224, 224), # img target size
|
||||
mean=mean,
|
||||
swapRB=True, # BGR -> RGB
|
||||
crop=True # center crop
|
||||
)
|
||||
# 3. divide by std
|
||||
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)
|
||||
return input_blob
|
||||
|
||||
|
||||
def get_imagenet_labels(labels_path):
|
||||
with open(labels_path) as f:
|
||||
imagenet_labels = [line.strip() for line in f.readlines()]
|
||||
return imagenet_labels
|
||||
|
||||
|
||||
def get_opencv_dnn_prediction(opencv_net, preproc_img, imagenet_labels):
|
||||
# set OpenCV DNN input
|
||||
opencv_net.setInput(preproc_img)
|
||||
|
||||
# OpenCV DNN inference
|
||||
out = opencv_net.forward()
|
||||
print("OpenCV DNN prediction: \n")
|
||||
print("* shape: ", out.shape)
|
||||
|
||||
# get the predicted class ID
|
||||
imagenet_class_id = np.argmax(out)
|
||||
|
||||
# get confidence
|
||||
confidence = out[0][imagenet_class_id]
|
||||
print("* class ID: {}, label: {}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
|
||||
print("* confidence: {:.4f}".format(confidence))
|
||||
|
||||
|
||||
def get_pytorch_dnn_prediction(original_net, preproc_img, imagenet_labels):
|
||||
original_net.eval()
|
||||
preproc_img = torch.FloatTensor(preproc_img)
|
||||
|
||||
# inference
|
||||
with torch.no_grad():
|
||||
out = original_net(preproc_img)
|
||||
|
||||
print("\nPyTorch model prediction: \n")
|
||||
print("* shape: ", out.shape)
|
||||
|
||||
# get the predicted class ID
|
||||
imagenet_class_id = torch.argmax(out, axis=1).item()
|
||||
print("* class ID: {}, label: {}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
|
||||
|
||||
# get confidence
|
||||
confidence = out[0][imagenet_class_id]
|
||||
print("* confidence: {:.4f}".format(confidence.item()))
|
||||
|
||||
|
||||
def main():
|
||||
# initialize PyTorch ResNet-50 model
|
||||
original_model = models.resnet50(pretrained=True)
|
||||
|
||||
# get the path to the converted into ONNX PyTorch model
|
||||
full_model_path = get_pytorch_onnx_model(original_model)
|
||||
|
||||
# read converted .onnx model with OpenCV API
|
||||
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
|
||||
print("OpenCV model was successfully read. Layer IDs: \n", opencv_net.getLayerNames())
|
||||
|
||||
# get preprocessed image
|
||||
input_img = get_preprocessed_img("../data/squirrel_cls.jpg")
|
||||
|
||||
# get ImageNet labels
|
||||
imagenet_labels = get_imagenet_labels("../data/dnn/classification_classes_ILSVRC2012.txt")
|
||||
|
||||
# obtain OpenCV DNN predictions
|
||||
get_opencv_dnn_prediction(opencv_net, input_img, imagenet_labels)
|
||||
|
||||
# obtain original PyTorch ResNet50 predictions
|
||||
get_pytorch_dnn_prediction(original_model, input_img, imagenet_labels)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+50
@@ -0,0 +1,50 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.onnx
|
||||
from torch.autograd import Variable
|
||||
from torchvision import models
|
||||
|
||||
|
||||
def get_pytorch_onnx_model(original_model):
|
||||
# define the directory for further converted model save
|
||||
onnx_model_path = "models"
|
||||
# define the name of further converted model
|
||||
onnx_model_name = "resnet50.onnx"
|
||||
|
||||
# create directory for further converted model
|
||||
os.makedirs(onnx_model_path, exist_ok=True)
|
||||
|
||||
# get full path to the converted model
|
||||
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
|
||||
|
||||
# generate model input
|
||||
generated_input = Variable(
|
||||
torch.randn(1, 3, 224, 224)
|
||||
)
|
||||
|
||||
# model export into ONNX format
|
||||
torch.onnx.export(
|
||||
original_model,
|
||||
generated_input,
|
||||
full_model_path,
|
||||
verbose=True,
|
||||
input_names=["input"],
|
||||
output_names=["output"],
|
||||
opset_version=11
|
||||
)
|
||||
|
||||
return full_model_path
|
||||
|
||||
|
||||
def main():
|
||||
# initialize PyTorch ResNet-50 model
|
||||
original_model = models.resnet50(pretrained=True)
|
||||
|
||||
# get the path to the converted into ONNX PyTorch model
|
||||
full_model_path = get_pytorch_onnx_model(original_model)
|
||||
print("PyTorch ResNet-50 model was successfully converted: ", full_model_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,98 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import torch.onnx
|
||||
from torch.autograd import Variable
|
||||
|
||||
from ..common.abstract_model import AbstractModel, Framework
|
||||
from ..common.utils import DNN_LIB, get_full_model_path
|
||||
|
||||
CURRENT_LIB = "PyTorch"
|
||||
MODEL_FORMAT = ".onnx"
|
||||
|
||||
|
||||
class PyTorchModelPreparer(AbstractModel):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
height,
|
||||
width,
|
||||
model_name="default",
|
||||
original_model=object,
|
||||
batch_size=1,
|
||||
default_input_name="input",
|
||||
default_output_name="output"
|
||||
):
|
||||
self._height = height
|
||||
self._width = width
|
||||
self._model_name = model_name
|
||||
self._original_model = original_model
|
||||
self._batch_size = batch_size
|
||||
self._default_input_name = default_input_name
|
||||
self._default_output_name = default_output_name
|
||||
|
||||
self.model_path = self._set_model_path()
|
||||
self._dnn_model = self._set_dnn_model()
|
||||
|
||||
def _set_dnn_model(self):
|
||||
generated_input = Variable(torch.randn(
|
||||
self._batch_size, 3, self._height, self._width)
|
||||
)
|
||||
os.makedirs(self.model_path["path"], exist_ok=True)
|
||||
torch.onnx.export(
|
||||
self._original_model,
|
||||
generated_input,
|
||||
self.model_path["full_path"],
|
||||
verbose=True,
|
||||
input_names=[self._default_input_name],
|
||||
output_names=[self._default_output_name],
|
||||
opset_version=11
|
||||
)
|
||||
|
||||
return cv2.dnn.readNetFromONNX(self.model_path["full_path"])
|
||||
|
||||
def _set_model_path(self):
|
||||
model_to_save = self._model_name + MODEL_FORMAT
|
||||
return get_full_model_path(CURRENT_LIB.lower(), model_to_save)
|
||||
|
||||
def get_prepared_models(self):
|
||||
return {
|
||||
CURRENT_LIB + " " + self._model_name: self._original_model,
|
||||
DNN_LIB + " " + self._model_name: self._dnn_model
|
||||
}
|
||||
|
||||
|
||||
class PyTorchModelProcessor(Framework):
|
||||
def __init__(self, prepared_model, model_name):
|
||||
self._prepared_model = prepared_model
|
||||
self._name = model_name
|
||||
|
||||
def get_output(self, input_blob):
|
||||
tensor = torch.FloatTensor(input_blob)
|
||||
self._prepared_model.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
model_out = self._prepared_model(tensor)
|
||||
|
||||
# segmentation case
|
||||
if len(model_out) == 2:
|
||||
model_out = model_out['out']
|
||||
|
||||
out = model_out.detach().numpy()
|
||||
return out
|
||||
|
||||
def get_name(self):
|
||||
return self._name
|
||||
|
||||
|
||||
class PyTorchDnnModelProcessor(Framework):
|
||||
def __init__(self, prepared_dnn_model, model_name):
|
||||
self._prepared_dnn_model = prepared_dnn_model
|
||||
self._name = model_name
|
||||
|
||||
def get_output(self, input_blob):
|
||||
self._prepared_dnn_model.setInput(input_blob, '')
|
||||
return self._prepared_dnn_model.forward()
|
||||
|
||||
def get_name(self):
|
||||
return self._name
|
||||
@@ -0,0 +1,15 @@
|
||||
# Python 3.7.5
|
||||
onnx>=1.7.0
|
||||
numpy>=1.19.1
|
||||
|
||||
torch>=1.5.1
|
||||
torchvision>=0.6.1
|
||||
|
||||
tensorflow>=2.1.0
|
||||
tensorflow-gpu>=2.1.0
|
||||
|
||||
paddlepaddle>=2.0.0
|
||||
paddlepaddle-gpu>=2.0.0
|
||||
paddlehub>=2.1.0
|
||||
paddle2onnx>=0.5.1
|
||||
paddleseg>=2.0.0
|
||||
@@ -0,0 +1,104 @@
|
||||
from tensorflow.keras.applications import (
|
||||
VGG16, vgg16,
|
||||
VGG19, vgg19,
|
||||
|
||||
ResNet50, resnet,
|
||||
ResNet101,
|
||||
ResNet152,
|
||||
|
||||
DenseNet121, densenet,
|
||||
DenseNet169,
|
||||
DenseNet201,
|
||||
|
||||
InceptionResNetV2, inception_resnet_v2,
|
||||
InceptionV3, inception_v3,
|
||||
|
||||
MobileNet, mobilenet,
|
||||
MobileNetV2, mobilenet_v2,
|
||||
|
||||
NASNetLarge, nasnet,
|
||||
NASNetMobile,
|
||||
|
||||
Xception, xception
|
||||
)
|
||||
|
||||
from ..tf_model import TFModelPreparer
|
||||
from ..tf_model import (
|
||||
TFModelProcessor,
|
||||
TFDnnModelProcessor
|
||||
)
|
||||
from ...common.evaluation.classification.cls_data_fetcher import TFPreprocessedFetch
|
||||
from ...common.test.cls_model_test_pipeline import ClsModelTestPipeline
|
||||
from ...common.test.configs.default_preprocess_config import (
|
||||
tf_input_blob,
|
||||
pytorch_input_blob,
|
||||
tf_model_blob_caffe_mode
|
||||
)
|
||||
from ...common.utils import set_tf_env, create_extended_parser
|
||||
|
||||
model_dict = {
|
||||
"vgg16": [VGG16, vgg16, tf_model_blob_caffe_mode],
|
||||
"vgg19": [VGG19, vgg19, tf_model_blob_caffe_mode],
|
||||
|
||||
"resnet50": [ResNet50, resnet, tf_model_blob_caffe_mode],
|
||||
"resnet101": [ResNet101, resnet, tf_model_blob_caffe_mode],
|
||||
"resnet152": [ResNet152, resnet, tf_model_blob_caffe_mode],
|
||||
|
||||
"densenet121": [DenseNet121, densenet, pytorch_input_blob],
|
||||
"densenet169": [DenseNet169, densenet, pytorch_input_blob],
|
||||
"densenet201": [DenseNet201, densenet, pytorch_input_blob],
|
||||
|
||||
"inceptionresnetv2": [InceptionResNetV2, inception_resnet_v2, tf_input_blob],
|
||||
"inceptionv3": [InceptionV3, inception_v3, tf_input_blob],
|
||||
|
||||
"mobilenet": [MobileNet, mobilenet, tf_input_blob],
|
||||
"mobilenetv2": [MobileNetV2, mobilenet_v2, tf_input_blob],
|
||||
|
||||
"nasnetlarge": [NASNetLarge, nasnet, tf_input_blob],
|
||||
"nasnetmobile": [NASNetMobile, nasnet, tf_input_blob],
|
||||
|
||||
"xception": [Xception, xception, tf_input_blob]
|
||||
}
|
||||
|
||||
CNN_CLASS_ID = 0
|
||||
CNN_UTILS_ID = 1
|
||||
DEFAULT_BLOB_PARAMS_ID = 2
|
||||
|
||||
|
||||
class TFClsModel(TFModelPreparer):
|
||||
def __init__(self, model_name, original_model):
|
||||
super(TFClsModel, self).__init__(model_name, original_model)
|
||||
|
||||
|
||||
def main():
|
||||
set_tf_env()
|
||||
|
||||
parser = create_extended_parser(list(model_dict.keys()))
|
||||
cmd_args = parser.parse_args()
|
||||
|
||||
model_name = cmd_args.model_name
|
||||
model_name_val = model_dict[model_name]
|
||||
|
||||
cls_model = TFClsModel(
|
||||
model_name=model_name,
|
||||
original_model=model_name_val[CNN_CLASS_ID](
|
||||
include_top=True,
|
||||
weights="imagenet"
|
||||
)
|
||||
)
|
||||
|
||||
tf_cls_pipeline = ClsModelTestPipeline(
|
||||
network_model=cls_model,
|
||||
model_processor=TFModelProcessor,
|
||||
dnn_model_processor=TFDnnModelProcessor,
|
||||
data_fetcher=TFPreprocessedFetch,
|
||||
img_processor=model_name_val[CNN_UTILS_ID].preprocess_input,
|
||||
cls_args_parser=parser,
|
||||
default_input_blob_preproc=model_name_val[DEFAULT_BLOB_PARAMS_ID]
|
||||
)
|
||||
|
||||
tf_cls_pipeline.init_test_pipeline()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,142 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.applications import MobileNet
|
||||
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
|
||||
|
||||
from ...common.utils import set_tf_env
|
||||
|
||||
|
||||
def get_tf_model_proto(tf_model):
|
||||
# define the directory for .pb model
|
||||
pb_model_path = "models"
|
||||
|
||||
# define the name of .pb model
|
||||
pb_model_name = "mobilenet.pb"
|
||||
|
||||
# create directory for further converted model
|
||||
os.makedirs(pb_model_path, exist_ok=True)
|
||||
|
||||
# get model TF graph
|
||||
tf_model_graph = tf.function(lambda x: tf_model(x))
|
||||
|
||||
# get concrete function
|
||||
tf_model_graph = tf_model_graph.get_concrete_function(
|
||||
tf.TensorSpec(tf_model.inputs[0].shape, tf_model.inputs[0].dtype))
|
||||
|
||||
# obtain frozen concrete function
|
||||
frozen_tf_func = convert_variables_to_constants_v2(tf_model_graph)
|
||||
# get frozen graph
|
||||
frozen_tf_func.graph.as_graph_def()
|
||||
|
||||
# save full tf model
|
||||
tf.io.write_graph(graph_or_graph_def=frozen_tf_func.graph,
|
||||
logdir=pb_model_path,
|
||||
name=pb_model_name,
|
||||
as_text=False)
|
||||
|
||||
return os.path.join(pb_model_path, pb_model_name)
|
||||
|
||||
|
||||
def get_preprocessed_img(img_path):
|
||||
# read the image
|
||||
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
|
||||
input_img = input_img.astype(np.float32)
|
||||
|
||||
# define preprocess parameters
|
||||
mean = np.array([1.0, 1.0, 1.0]) * 127.5
|
||||
scale = 1 / 127.5
|
||||
|
||||
# prepare input blob to fit the model input:
|
||||
# 1. subtract mean
|
||||
# 2. scale to set pixel values from 0 to 1
|
||||
input_blob = cv2.dnn.blobFromImage(
|
||||
image=input_img,
|
||||
scalefactor=scale,
|
||||
size=(224, 224), # img target size
|
||||
mean=mean,
|
||||
swapRB=True, # BGR -> RGB
|
||||
crop=True # center crop
|
||||
)
|
||||
print("Input blob shape: {}\n".format(input_blob.shape))
|
||||
|
||||
return input_blob
|
||||
|
||||
|
||||
def get_imagenet_labels(labels_path):
|
||||
with open(labels_path) as f:
|
||||
imagenet_labels = [line.strip() for line in f.readlines()]
|
||||
return imagenet_labels
|
||||
|
||||
|
||||
def get_opencv_dnn_prediction(opencv_net, preproc_img, imagenet_labels):
|
||||
# set OpenCV DNN input
|
||||
opencv_net.setInput(preproc_img)
|
||||
|
||||
# OpenCV DNN inference
|
||||
out = opencv_net.forward()
|
||||
print("OpenCV DNN prediction: \n")
|
||||
print("* shape: ", out.shape)
|
||||
|
||||
# get the predicted class ID
|
||||
imagenet_class_id = np.argmax(out)
|
||||
|
||||
# get confidence
|
||||
confidence = out[0][imagenet_class_id]
|
||||
print("* class ID: {}, label: {}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
|
||||
print("* confidence: {:.4f}\n".format(confidence))
|
||||
|
||||
|
||||
def get_tf_dnn_prediction(original_net, preproc_img, imagenet_labels):
|
||||
# inference
|
||||
preproc_img = preproc_img.transpose(0, 2, 3, 1)
|
||||
print("TF input blob shape: {}\n".format(preproc_img.shape))
|
||||
|
||||
out = original_net(preproc_img)
|
||||
|
||||
print("\nTensorFlow model prediction: \n")
|
||||
print("* shape: ", out.shape)
|
||||
|
||||
# get the predicted class ID
|
||||
imagenet_class_id = np.argmax(out)
|
||||
print("* class ID: {}, label: {}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
|
||||
|
||||
# get confidence
|
||||
confidence = out[0][imagenet_class_id]
|
||||
print("* confidence: {:.4f}".format(confidence))
|
||||
|
||||
|
||||
def main():
|
||||
# configure TF launching
|
||||
set_tf_env()
|
||||
|
||||
# initialize TF MobileNet model
|
||||
original_tf_model = MobileNet(
|
||||
include_top=True,
|
||||
weights="imagenet"
|
||||
)
|
||||
|
||||
# get TF frozen graph path
|
||||
full_pb_path = get_tf_model_proto(original_tf_model)
|
||||
|
||||
# read frozen graph with OpenCV API
|
||||
opencv_net = cv2.dnn.readNetFromTensorflow(full_pb_path)
|
||||
print("OpenCV model was successfully read. Model layers: \n", opencv_net.getLayerNames())
|
||||
|
||||
# get preprocessed image
|
||||
input_img = get_preprocessed_img("../data/squirrel_cls.jpg")
|
||||
|
||||
# get ImageNet labels
|
||||
imagenet_labels = get_imagenet_labels("../data/dnn/classification_classes_ILSVRC2012.txt")
|
||||
|
||||
# obtain OpenCV DNN predictions
|
||||
get_opencv_dnn_prediction(opencv_net, input_img, imagenet_labels)
|
||||
|
||||
# obtain TF model predictions
|
||||
get_tf_dnn_prediction(original_tf_model, input_img, imagenet_labels)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,45 @@
|
||||
import os
|
||||
import tarfile
|
||||
import urllib
|
||||
|
||||
DETECTION_MODELS_URL = 'http://download.tensorflow.org/models/object_detection/'
|
||||
|
||||
|
||||
def extract_tf_frozen_graph(model_name, extracted_model_path):
|
||||
# define model archive name
|
||||
tf_model_tar = model_name + '.tar.gz'
|
||||
# define link to retrieve model archive
|
||||
model_link = DETECTION_MODELS_URL + tf_model_tar
|
||||
|
||||
tf_frozen_graph_name = 'frozen_inference_graph'
|
||||
|
||||
try:
|
||||
urllib.request.urlretrieve(model_link, tf_model_tar)
|
||||
except Exception:
|
||||
print("TF {} was not retrieved: {}".format(model_name, model_link))
|
||||
return
|
||||
|
||||
print("TF {} was retrieved.".format(model_name))
|
||||
|
||||
tf_model_tar = tarfile.open(tf_model_tar)
|
||||
frozen_graph_path = ""
|
||||
|
||||
for model_tar_elem in tf_model_tar.getmembers():
|
||||
if tf_frozen_graph_name in os.path.basename(model_tar_elem.name):
|
||||
tf_model_tar.extract(model_tar_elem, extracted_model_path)
|
||||
frozen_graph_path = os.path.join(extracted_model_path, model_tar_elem.name)
|
||||
break
|
||||
tf_model_tar.close()
|
||||
|
||||
return frozen_graph_path
|
||||
|
||||
|
||||
def main():
|
||||
tf_model_name = 'ssd_mobilenet_v1_coco_2017_11_17'
|
||||
graph_extraction_dir = "./"
|
||||
frozen_graph_path = extract_tf_frozen_graph(tf_model_name, graph_extraction_dir)
|
||||
print("Frozen graph path for {}: {}".format(tf_model_name, frozen_graph_path))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,112 @@
|
||||
import cv2
|
||||
import tensorflow as tf
|
||||
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
|
||||
|
||||
from ..common.abstract_model import AbstractModel, Framework
|
||||
from ..common.utils import DNN_LIB, get_full_model_path
|
||||
|
||||
CURRENT_LIB = "TF"
|
||||
MODEL_FORMAT = ".pb"
|
||||
|
||||
|
||||
class TFModelPreparer(AbstractModel):
|
||||
""" Class for the preparation of the TF models: original and converted OpenCV Net.
|
||||
|
||||
Args:
|
||||
model_name: TF model name
|
||||
original_model: TF configured model object or session
|
||||
is_ready_graph: indicates whether ready .pb file already exists
|
||||
tf_model_graph_path: path to the existing frozen TF graph
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name="default",
|
||||
original_model=None,
|
||||
is_ready_graph=False,
|
||||
tf_model_graph_path=""
|
||||
):
|
||||
self._model_name = model_name
|
||||
self._original_model = original_model
|
||||
self._model_to_save = ""
|
||||
|
||||
self._is_ready_to_transfer_graph = is_ready_graph
|
||||
self.model_path = self._set_model_path(tf_model_graph_path)
|
||||
self._dnn_model = self._set_dnn_model()
|
||||
|
||||
def _set_dnn_model(self):
|
||||
if not self._is_ready_to_transfer_graph:
|
||||
# get model TF graph
|
||||
tf_model_graph = tf.function(lambda x: self._original_model(x))
|
||||
|
||||
tf_model_graph = tf_model_graph.get_concrete_function(
|
||||
tf.TensorSpec(self._original_model.inputs[0].shape, self._original_model.inputs[0].dtype))
|
||||
|
||||
# obtain frozen concrete function
|
||||
frozen_tf_func = convert_variables_to_constants_v2(tf_model_graph)
|
||||
frozen_tf_func.graph.as_graph_def()
|
||||
|
||||
# save full TF model
|
||||
tf.io.write_graph(graph_or_graph_def=frozen_tf_func.graph,
|
||||
logdir=self.model_path["path"],
|
||||
name=self._model_to_save,
|
||||
as_text=False)
|
||||
|
||||
return cv2.dnn.readNetFromTensorflow(self.model_path["full_path"])
|
||||
|
||||
def _set_model_path(self, tf_pb_file_path):
|
||||
""" Method for setting model paths.
|
||||
|
||||
Args:
|
||||
tf_pb_file_path: path to the existing TF .pb
|
||||
|
||||
Returns:
|
||||
dictionary, where full_path key means saved model path and its full name.
|
||||
"""
|
||||
model_paths_dict = {
|
||||
"path": "",
|
||||
"full_path": tf_pb_file_path
|
||||
}
|
||||
|
||||
if not self._is_ready_to_transfer_graph:
|
||||
self._model_to_save = self._model_name + MODEL_FORMAT
|
||||
model_paths_dict = get_full_model_path(CURRENT_LIB.lower(), self._model_to_save)
|
||||
|
||||
return model_paths_dict
|
||||
|
||||
def get_prepared_models(self):
|
||||
original_lib_name = CURRENT_LIB + " " + self._model_name
|
||||
configured_model_dict = {
|
||||
original_lib_name: self._original_model,
|
||||
DNN_LIB + " " + self._model_name: self._dnn_model
|
||||
}
|
||||
return configured_model_dict
|
||||
|
||||
|
||||
class TFModelProcessor(Framework):
|
||||
def __init__(self, prepared_model, model_name):
|
||||
self._prepared_model = prepared_model
|
||||
self._name = model_name
|
||||
|
||||
def get_output(self, input_blob):
|
||||
assert len(input_blob.shape) == 4
|
||||
batch_tf = input_blob.transpose(0, 2, 3, 1)
|
||||
out = self._prepared_model(batch_tf)
|
||||
return out
|
||||
|
||||
def get_name(self):
|
||||
return CURRENT_LIB
|
||||
|
||||
|
||||
class TFDnnModelProcessor(Framework):
|
||||
def __init__(self, prepared_dnn_model, model_name):
|
||||
self._prepared_dnn_model = prepared_dnn_model
|
||||
self._name = model_name
|
||||
|
||||
def get_output(self, input_blob):
|
||||
self._prepared_dnn_model.setInput(input_blob)
|
||||
ret_val = self._prepared_dnn_model.forward()
|
||||
return ret_val
|
||||
|
||||
def get_name(self):
|
||||
return DNN_LIB
|
||||
@@ -0,0 +1,383 @@
|
||||
'''
|
||||
Helper module to download extra data from Internet
|
||||
'''
|
||||
from __future__ import print_function
|
||||
import os
|
||||
import sys
|
||||
import yaml
|
||||
import argparse
|
||||
import tarfile
|
||||
import platform
|
||||
import tempfile
|
||||
import hashlib
|
||||
import requests
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from urllib.request import Request, urlopen
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
__all__ = ["downloadFile"]
|
||||
|
||||
class HashMismatchException(Exception):
|
||||
def __init__(self, expected, actual):
|
||||
Exception.__init__(self)
|
||||
self.expected = expected
|
||||
self.actual = actual
|
||||
def __str__(self):
|
||||
return 'Hash mismatch: expected {} vs actual of {}'.format(self.expected, self.actual)
|
||||
|
||||
def getHashsumFromFile(filepath):
|
||||
sha = hashlib.sha1()
|
||||
if os.path.exists(filepath):
|
||||
print(' there is already a file with the same name')
|
||||
with open(filepath, 'rb') as f:
|
||||
while True:
|
||||
buf = f.read(10*1024*1024)
|
||||
if not buf:
|
||||
break
|
||||
sha.update(buf)
|
||||
hashsum = sha.hexdigest()
|
||||
return hashsum
|
||||
|
||||
def checkHashsum(expected_sha, filepath, silent=True):
|
||||
if not os.path.exists(filepath):
|
||||
print(f"{filepath} does not exist. Skipping hashsum matching")
|
||||
return False
|
||||
print(' expected SHA1: {}'.format(expected_sha))
|
||||
actual_sha = getHashsumFromFile(filepath)
|
||||
print(' actual SHA1:{}'.format(actual_sha))
|
||||
hashes_matched = expected_sha == actual_sha
|
||||
if not hashes_matched and not silent:
|
||||
raise HashMismatchException(expected_sha, actual_sha)
|
||||
return hashes_matched
|
||||
|
||||
def isArchive(filepath):
|
||||
return tarfile.is_tarfile(filepath)
|
||||
|
||||
class DownloadInstance:
|
||||
def __init__(self, **kwargs):
|
||||
self.name = kwargs.pop('name')
|
||||
self.filename = kwargs.pop('filename')
|
||||
self.loader = kwargs.pop('loader', None)
|
||||
self.save_dir = kwargs.pop('save_dir')
|
||||
self.sha = kwargs.pop('sha', None)
|
||||
|
||||
def __str__(self):
|
||||
return 'DownloadInstance <{}>'.format(self.name)
|
||||
|
||||
def get(self):
|
||||
print(" Working on " + self.name)
|
||||
print(" Getting file " + self.filename)
|
||||
if self.sha is None:
|
||||
print(' No expected hashsum provided, loading file')
|
||||
else:
|
||||
filepath = os.path.join(self.save_dir, self.sha, self.filename)
|
||||
if checkHashsum(self.sha, filepath):
|
||||
print(' hash match - file already exists, skipping')
|
||||
return filepath
|
||||
else:
|
||||
print(' hash didn\'t match, loading file')
|
||||
|
||||
if not os.path.exists(self.save_dir):
|
||||
print(' creating directory: ' + self.save_dir)
|
||||
os.makedirs(self.save_dir)
|
||||
|
||||
|
||||
print(' hash check failed - loading')
|
||||
assert self.loader
|
||||
try:
|
||||
self.loader.load(self.filename, self.sha, self.save_dir)
|
||||
print(' done')
|
||||
print(' file {}'.format(self.filename))
|
||||
if self.sha is None:
|
||||
download_path = os.path.join(self.save_dir, self.filename)
|
||||
self.sha = getHashsumFromFile(download_path)
|
||||
new_dir = os.path.join(self.save_dir, self.sha)
|
||||
|
||||
if not os.path.exists(new_dir):
|
||||
os.makedirs(new_dir)
|
||||
filepath = os.path.join(new_dir, self.filename)
|
||||
if not (os.path.exists(filepath)):
|
||||
shutil.move(download_path, new_dir)
|
||||
print(' No expected hashsum provided, actual SHA is {}'.format(self.sha))
|
||||
else:
|
||||
checkHashsum(self.sha, filepath, silent=False)
|
||||
except Exception as e:
|
||||
print(" There was some problem with loading file {} for {}".format(self.filename, self.name))
|
||||
print(" Exception: {}".format(e))
|
||||
return
|
||||
|
||||
print(" Finished " + self.name)
|
||||
return filepath
|
||||
|
||||
class Loader(object):
|
||||
MB = 1024*1024
|
||||
BUFSIZE = 10*MB
|
||||
def __init__(self, download_name, download_sha, archive_member = None):
|
||||
self.download_name = download_name
|
||||
self.download_sha = download_sha
|
||||
self.archive_member = archive_member
|
||||
|
||||
def load(self, requested_file, sha, save_dir):
|
||||
if self.download_sha is None:
|
||||
download_dir = save_dir
|
||||
else:
|
||||
# create a new folder in save_dir to avoid possible name conflicts
|
||||
download_dir = os.path.join(save_dir, self.download_sha)
|
||||
if not os.path.exists(download_dir):
|
||||
os.makedirs(download_dir)
|
||||
download_path = os.path.join(download_dir, self.download_name)
|
||||
print(" Preparing to download file " + self.download_name)
|
||||
if checkHashsum(self.download_sha, download_path):
|
||||
print(' hash match - file already exists, no need to download')
|
||||
else:
|
||||
filesize = self.download(download_path)
|
||||
print(' Downloaded {} with size {} Mb'.format(self.download_name, filesize/self.MB))
|
||||
if self.download_sha is not None:
|
||||
checkHashsum(self.download_sha, download_path, silent=False)
|
||||
if self.download_name == requested_file:
|
||||
return
|
||||
else:
|
||||
if isArchive(download_path):
|
||||
if sha is not None:
|
||||
extract_dir = os.path.join(save_dir, sha)
|
||||
else:
|
||||
extract_dir = save_dir
|
||||
if not os.path.exists(extract_dir):
|
||||
os.makedirs(extract_dir)
|
||||
self.extract(requested_file, download_path, extract_dir)
|
||||
else:
|
||||
raise Exception("Downloaded file has different name")
|
||||
|
||||
def download(self, filepath):
|
||||
print("Warning: download is not implemented, this is a base class")
|
||||
return 0
|
||||
|
||||
def extract(self, requested_file, archive_path, save_dir):
|
||||
filepath = os.path.join(save_dir, requested_file)
|
||||
try:
|
||||
with tarfile.open(archive_path) as f:
|
||||
if self.archive_member is None:
|
||||
pathDict = dict((os.path.split(elem)[1], os.path.split(elem)[0]) for elem in f.getnames())
|
||||
self.archive_member = pathDict[requested_file]
|
||||
if self.archive_member == "":
|
||||
self.archive_member = requested_file
|
||||
assert self.archive_member in f.getnames()
|
||||
self.save(filepath, f.extractfile(self.archive_member))
|
||||
except Exception as e:
|
||||
print(' catch {}'.format(e))
|
||||
|
||||
def save(self, filepath, r):
|
||||
with open(filepath, 'wb') as f:
|
||||
print(' progress ', end="")
|
||||
sys.stdout.flush()
|
||||
while True:
|
||||
buf = r.read(self.BUFSIZE)
|
||||
if not buf:
|
||||
break
|
||||
f.write(buf)
|
||||
print('>', end="")
|
||||
sys.stdout.flush()
|
||||
|
||||
class URLLoader(Loader):
|
||||
def __init__(self, download_name, download_sha, url, archive_member = None):
|
||||
super(URLLoader, self).__init__(download_name, download_sha, archive_member)
|
||||
self.download_name = download_name
|
||||
self.download_sha = download_sha
|
||||
self.url = url
|
||||
|
||||
def download(self, filepath):
|
||||
headers = {'User-Agent': 'Wget/1.20.3'}
|
||||
req = Request(self.url, headers=headers)
|
||||
with urlopen(req, timeout=60) as r:
|
||||
self.printRequest(r)
|
||||
self.save(filepath, r)
|
||||
return os.path.getsize(filepath)
|
||||
|
||||
def printRequest(self, r):
|
||||
def getMB(r):
|
||||
d = dict(r.info())
|
||||
for c in ['content-length', 'Content-Length']:
|
||||
if c in d:
|
||||
return int(d[c]) / self.MB
|
||||
return '<unknown>'
|
||||
print(' {} {} [{} Mb]'.format(r.getcode(), r.msg, getMB(r)))
|
||||
|
||||
class GDriveLoader(Loader):
|
||||
BUFSIZE = 1024 * 1024
|
||||
PROGRESS_SIZE = 10 * 1024 * 1024
|
||||
def __init__(self, download_name, download_sha, gid, archive_member = None):
|
||||
super(GDriveLoader, self).__init__(download_name, download_sha, archive_member)
|
||||
self.download_name = download_name
|
||||
self.download_sha = download_sha
|
||||
self.gid = gid
|
||||
|
||||
def download(self, filepath):
|
||||
session = requests.Session() # re-use cookies
|
||||
|
||||
URL = "https://docs.google.com/uc?export=download"
|
||||
response = session.get(URL, params = { 'id' : self.gid }, stream = True)
|
||||
|
||||
def get_confirm_token(response): # in case of large files
|
||||
for key, value in response.cookies.items():
|
||||
if key.startswith('download_warning'):
|
||||
return value
|
||||
return None
|
||||
token = get_confirm_token(response)
|
||||
|
||||
if token:
|
||||
params = { 'id' : self.gid, 'confirm' : token }
|
||||
response = session.get(URL, params = params, stream = True)
|
||||
|
||||
sz = 0
|
||||
progress_sz = self.PROGRESS_SIZE
|
||||
with open(filepath, "wb") as f:
|
||||
for chunk in response.iter_content(self.BUFSIZE):
|
||||
if not chunk:
|
||||
continue # keep-alive
|
||||
|
||||
f.write(chunk)
|
||||
sz += len(chunk)
|
||||
if sz >= progress_sz:
|
||||
progress_sz += self.PROGRESS_SIZE
|
||||
print('>', end='')
|
||||
sys.stdout.flush()
|
||||
print('')
|
||||
return sz
|
||||
|
||||
def produceDownloadInstance(instance_name, filename, sha, url, save_dir, download_name=None, download_sha=None, archive_member=None):
|
||||
spec_param = url
|
||||
loader = URLLoader
|
||||
if download_name is None:
|
||||
download_name = filename
|
||||
if download_sha is None:
|
||||
download_sha = sha
|
||||
if "drive.google.com" in url:
|
||||
token = ""
|
||||
token_part = url.rsplit('/', 1)[-1]
|
||||
if "&id=" not in token_part:
|
||||
token_part = url.rsplit('/', 1)[-2]
|
||||
for param in token_part.split("&"):
|
||||
if param.startswith("id="):
|
||||
token = param[3:]
|
||||
if token:
|
||||
loader = GDriveLoader
|
||||
spec_param = token
|
||||
else:
|
||||
print("Warning: possibly wrong Google Drive link")
|
||||
return DownloadInstance(
|
||||
name=instance_name,
|
||||
filename=filename,
|
||||
sha=sha,
|
||||
save_dir=save_dir,
|
||||
loader=loader(download_name, download_sha, spec_param, archive_member)
|
||||
)
|
||||
|
||||
def getSaveDir():
|
||||
env_path = os.environ.get("OPENCV_DOWNLOAD_DATA_PATH", None)
|
||||
if env_path:
|
||||
save_dir = env_path
|
||||
else:
|
||||
# TODO reuse binding function cv2.utils.fs.getCacheDirectory when issue #19011 is fixed
|
||||
if platform.system() == "Darwin":
|
||||
#On Apple devices
|
||||
temp_env = os.environ.get("TMPDIR", None)
|
||||
if temp_env is None or not os.path.isdir(temp_env):
|
||||
temp_dir = Path("/tmp")
|
||||
print("Using world accessible cache directory. This may be not secure: ", temp_dir)
|
||||
else:
|
||||
temp_dir = temp_env
|
||||
elif platform.system() == "Windows":
|
||||
temp_dir = tempfile.gettempdir()
|
||||
else:
|
||||
xdg_cache_env = os.environ.get("XDG_CACHE_HOME", None)
|
||||
if (xdg_cache_env and xdg_cache_env[0] and os.path.isdir(xdg_cache_env)):
|
||||
temp_dir = xdg_cache_env
|
||||
else:
|
||||
home_env = os.environ.get("HOME", None)
|
||||
if (home_env and home_env[0] and os.path.isdir(home_env)):
|
||||
home_path = os.path.join(home_env, ".cache/")
|
||||
if os.path.isdir(home_path):
|
||||
temp_dir = home_path
|
||||
else:
|
||||
temp_dir = tempfile.gettempdir()
|
||||
print("Using world accessible cache directory. This may be not secure: ", temp_dir)
|
||||
|
||||
save_dir = os.path.join(temp_dir, "downloads")
|
||||
if not os.path.exists(save_dir):
|
||||
os.makedirs(save_dir)
|
||||
return save_dir
|
||||
|
||||
def downloadFile(url, sha=None, save_dir=None, filename=None):
|
||||
if save_dir is None:
|
||||
save_dir = getSaveDir()
|
||||
if filename is None:
|
||||
filename = "download_" + datetime.now().__str__()
|
||||
name = filename
|
||||
return produceDownloadInstance(name, filename, sha, url, save_dir).get()
|
||||
|
||||
def parseMetalinkFile(metalink_filepath, save_dir):
|
||||
NS = {'ml': 'urn:ietf:params:xml:ns:metalink'}
|
||||
models = []
|
||||
for file_elem in ET.parse(metalink_filepath).getroot().findall('ml:file', NS):
|
||||
url = file_elem.find('ml:url', NS).text
|
||||
fname = file_elem.attrib['name']
|
||||
name = file_elem.find('ml:identity', NS).text
|
||||
hash_sum = file_elem.find('ml:hash', NS).text
|
||||
models.append(produceDownloadInstance(name, fname, hash_sum, url, save_dir))
|
||||
return models
|
||||
|
||||
def parseYAMLFile(yaml_filepath, save_dir, model_name):
|
||||
models = []
|
||||
with open(yaml_filepath, 'r') as stream:
|
||||
data_loaded = yaml.safe_load(stream)
|
||||
for name, params in data_loaded.items():
|
||||
if model_name != "" and name != model_name:
|
||||
continue
|
||||
for key in params.keys():
|
||||
if key.endswith("load_info"):
|
||||
prefix = key[:-len('load_info')]
|
||||
load_info = params.get(prefix+"load_info", None)
|
||||
if load_info:
|
||||
print(prefix)
|
||||
if prefix == "config_":
|
||||
fname = os.path.basename(params.get("config"))
|
||||
hash_sum = load_info.get("sha1")
|
||||
url = load_info.get("url")
|
||||
download_sha = load_info.get("download_sha")
|
||||
download_name = load_info.get("download_name")
|
||||
archive_member = load_info.get("member")
|
||||
models.append(produceDownloadInstance(name, fname, hash_sum, url, save_dir,
|
||||
download_name=download_name, download_sha=download_sha, archive_member=archive_member))
|
||||
else:
|
||||
fname = os.path.basename(params.get(prefix+"model"))
|
||||
hash_sum = load_info.get(prefix+"sha1")
|
||||
url = load_info.get(prefix+"url")
|
||||
download_sha = load_info.get(prefix+"download_sha")
|
||||
download_name = load_info.get(prefix+"download_name")
|
||||
archive_member = load_info.get(prefix+"member")
|
||||
models.append(produceDownloadInstance(name, fname, hash_sum, url, save_dir,
|
||||
download_name=download_name, download_sha=download_sha, archive_member=archive_member))
|
||||
|
||||
return models
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='This is a utility script for downloading DNN models for samples.')
|
||||
|
||||
parser.add_argument('--save_dir', action="store", default=os.getcwd(),
|
||||
help='Path to the directory to store downloaded files')
|
||||
parser.add_argument('model_name', type=str, default="", nargs='?', action="store",
|
||||
help='name of the model to download')
|
||||
args = parser.parse_args()
|
||||
models = []
|
||||
save_dir = args.save_dir
|
||||
selected_model_name = args.model_name
|
||||
models.extend(parseMetalinkFile('face_detector/weights.meta4', save_dir))
|
||||
models.extend(parseYAMLFile('models.yml', save_dir, selected_model_name))
|
||||
for m in models:
|
||||
print(m)
|
||||
if selected_model_name and not m.name.startswith(selected_model_name):
|
||||
continue
|
||||
print('Model: ' + selected_model_name)
|
||||
m.get()
|
||||
@@ -0,0 +1,251 @@
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "common.hpp"
|
||||
// Define namespace to simplify code
|
||||
using namespace cv;
|
||||
using namespace cv::dnn;
|
||||
using namespace std;
|
||||
|
||||
int threshold1 = 0;
|
||||
int threshold2 = 50;
|
||||
int blurAmount = 5;
|
||||
|
||||
// Function to apply sigmoid activation
|
||||
static void sigmoid(Mat& input) {
|
||||
exp(-input, input); // e^-input
|
||||
input = 1.0 / (1.0 + input); // 1 / (1 + e^-input)
|
||||
}
|
||||
|
||||
static void applyCanny(const Mat& image, Mat& result) {
|
||||
Mat gray;
|
||||
cvtColor(image, gray, COLOR_BGR2GRAY);
|
||||
Canny(gray, result, threshold1, threshold2);
|
||||
}
|
||||
|
||||
// Load Model
|
||||
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 setupCannyWindow(){
|
||||
destroyWindow("Output");
|
||||
namedWindow("Output", WINDOW_AUTOSIZE);
|
||||
moveWindow("Output", 200, 50);
|
||||
|
||||
createTrackbar("thrs1", "Output", &threshold1, 255, nullptr);
|
||||
createTrackbar("thrs2", "Output", &threshold2, 255, nullptr);
|
||||
createTrackbar("blur", "Output", &blurAmount, 20, nullptr);
|
||||
}
|
||||
|
||||
// Function to process the neural network output to generate edge maps
|
||||
static pair<Mat, Mat> postProcess(const vector<Mat>& output, int height, int width) {
|
||||
vector<Mat> preds;
|
||||
preds.reserve(output.size());
|
||||
for (const Mat &p : output) {
|
||||
Mat img;
|
||||
// Correctly handle 4D tensor assuming it's always in the format [1, 1, height, width]
|
||||
Mat processed;
|
||||
if (p.dims == 4 && p.size[0] == 1 && p.size[1] == 1) {
|
||||
// Use only the spatial dimensions
|
||||
processed = p.reshape(0, {p.size[2], p.size[3]});
|
||||
} else {
|
||||
processed = p.clone();
|
||||
}
|
||||
sigmoid(processed);
|
||||
normalize(processed, img, 0, 255, NORM_MINMAX, CV_8U);
|
||||
resize(img, img, Size(width, height)); // Resize to the original size
|
||||
preds.push_back(img);
|
||||
}
|
||||
Mat fuse = preds.back(); // Last element as the fused result
|
||||
// Calculate the average of the predictions
|
||||
Mat ave = Mat::zeros(height, width, CV_32F);
|
||||
for (Mat &pred : preds) {
|
||||
Mat temp;
|
||||
pred.convertTo(temp, CV_32F);
|
||||
ave += temp;
|
||||
}
|
||||
ave /= static_cast<float>(preds.size());
|
||||
ave.convertTo(ave, CV_8U);
|
||||
return {fuse, ave}; // Return both fused and average edge maps
|
||||
}
|
||||
|
||||
static void applyDexined(Net &net, const Mat &image, Mat &result) {
|
||||
int originalWidth = image.cols;
|
||||
int originalHeight = image.rows;
|
||||
vector<Mat> outputs;
|
||||
net.forward(outputs);
|
||||
pair<Mat, Mat> res = postProcess(outputs, originalHeight, originalWidth);
|
||||
result = res.first; // or res.second for average edge map
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
const string about =
|
||||
"This sample demonstrates edge detection with dexined and canny edge detection techniques.\n\n"
|
||||
"To run with canny:\n"
|
||||
"\t ./example_dnn_edge_detection --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)\n"
|
||||
"With Dexined:\n"
|
||||
"\t ./example_dnn_edge_detection dexined --input=path/to/your/input/image/or/video\n\n"
|
||||
"For switching between deep learning based model(dexined) and canny edge detector, press space bar in case of video. In case of image, pass the argument --method for switching between dexined and canny.\n"
|
||||
"Model path can also be specified using --model argument. Download it using python download_models.py dexined from dnn samples directory\n\n";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ @alias | | 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 | | Path to input image or video file. Skip this argument to capture frames from a camera.}"
|
||||
"{ method | dexined | Choose method: dexined or canny. }"
|
||||
"{ model | | Path to the model file for using dexined. }";
|
||||
|
||||
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;
|
||||
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
if (parser.has("help"))
|
||||
{
|
||||
cout << about << endl;
|
||||
parser.printMessage();
|
||||
return -1;
|
||||
}
|
||||
|
||||
string modelName = parser.get<String>("@alias");
|
||||
string zooFile = parser.get<String>("zoo");
|
||||
|
||||
const char* path = getenv("OPENCV_SAMPLES_DATA_PATH");
|
||||
if ((path != NULL) || parser.has("@alias") || (parser.get<String>("model") != "")) {
|
||||
modelName = "dexined";
|
||||
zooFile = findFile(zooFile);
|
||||
}
|
||||
else{
|
||||
cout<<"[WARN] set the environment variables or pass path to dexined.onnx model file using --model and models.yml file using --zoo for using dexined based edge detector. Continuing with canny edge detector\n\n";
|
||||
}
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
int width = parser.get<int>("width");
|
||||
int height = parser.get<int>("height");
|
||||
float scale = parser.get<float>("scale");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
String backend = parser.get<String>("backend");
|
||||
String target = parser.get<String>("target");
|
||||
string method = parser.get<String>("method");
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
string model = findModel(parser.get<String>("model"), sha1);
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if (backend != "default" || target != "cpu"){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
parser.about(about);
|
||||
|
||||
VideoCapture cap;
|
||||
if (parser.has("input"))
|
||||
cap.open(samples::findFile(parser.get<String>("input")));
|
||||
else
|
||||
cap.open(0);
|
||||
|
||||
if (!cap.isOpened()) {
|
||||
cerr << "Error: Video could not be opened." << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
namedWindow("Input", WINDOW_AUTOSIZE);
|
||||
namedWindow("Output", WINDOW_AUTOSIZE);
|
||||
moveWindow("Output", 200, 0);
|
||||
Net net;
|
||||
Mat image;
|
||||
|
||||
if (model.empty()) {
|
||||
cout << "[WARN] Model file not provided, using canny instead. Pass model using --model=/path/to/dexined.onnx to use dexined model." << endl;
|
||||
method = "canny";
|
||||
}
|
||||
|
||||
if (method == "dexined") {
|
||||
loadModel(model, backend, target, net, engine);
|
||||
}
|
||||
else{
|
||||
Mat dummy = Mat::zeros(512, 512, CV_8UC3);
|
||||
setupCannyWindow();
|
||||
}
|
||||
cout<<"To switch between canny and dexined press space bar."<<endl;
|
||||
for (;;){
|
||||
cap >> image;
|
||||
if (image.empty())
|
||||
{
|
||||
cout << "Press any key to exit" << endl;
|
||||
waitKey();
|
||||
break;
|
||||
}
|
||||
|
||||
Mat result;
|
||||
int kernelSize = 2 * blurAmount + 1;
|
||||
Mat blurred;
|
||||
GaussianBlur(image, blurred, Size(kernelSize, kernelSize), 0);
|
||||
if (method == "dexined")
|
||||
{
|
||||
Mat blob = blobFromImage(blurred, scale, Size(width, height), mean, swapRB, false, CV_32F);
|
||||
net.setInput(blob);
|
||||
applyDexined(net, image, result);
|
||||
}
|
||||
else if (method == "canny")
|
||||
{
|
||||
applyCanny(blurred, result);
|
||||
}
|
||||
imshow("Input", image);
|
||||
imshow("Output", result);
|
||||
int key = waitKey(30);
|
||||
|
||||
if (key == ' ' && method == "canny")
|
||||
{
|
||||
if (!model.empty()){
|
||||
method = "dexined";
|
||||
if (net.empty())
|
||||
loadModel(model, backend, target, net, engine);
|
||||
destroyWindow("Output");
|
||||
namedWindow("Input", WINDOW_AUTOSIZE);
|
||||
namedWindow("Output", WINDOW_AUTOSIZE);
|
||||
moveWindow("Output", 200, 0);
|
||||
} else {
|
||||
cout << "[ERROR] Provide model file using --model to use dexined. Download model using python download_models.py dexined from dnn samples directory" << endl;
|
||||
}
|
||||
}
|
||||
else if (key == ' ' && method == "dexined")
|
||||
{
|
||||
method = "canny";
|
||||
setupCannyWindow();
|
||||
}
|
||||
else if (key == 27 || key == 'q')
|
||||
{ // Escape key to exit
|
||||
break;
|
||||
}
|
||||
}
|
||||
destroyAllWindows();
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
'''
|
||||
This sample demonstrates edge detection with dexined and canny edge detection techniques.
|
||||
For switching between deep learning based model(dexined) and canny edge detector, press space bar in case of video. In case of image, pass the argument --method for switching between dexined and canny.
|
||||
'''
|
||||
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import numpy as np
|
||||
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', help='Path to input image or video file. Skip this argument to capture frames from a camera.', default=0, required=False)
|
||||
parser.add_argument('--method', help='choose method: dexined or canny', default='canny', 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, 'edge_detection', 'dexined')
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='''
|
||||
To run:
|
||||
Canny:
|
||||
python edge_detection.py --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)
|
||||
Dexined:
|
||||
python edge_detection.py dexined --input=path/to/your/input/image/or/video
|
||||
|
||||
"In case of video input, for switching between deep learning based model (Dexined) and Canny edge detector, press space bar. Pass as argument in case of image input."
|
||||
|
||||
Model path can also be specified using --model argument
|
||||
''', formatter_class=argparse.RawTextHelpFormatter)
|
||||
return parser.parse_args(func_args)
|
||||
|
||||
threshold1 = 0
|
||||
threshold2 = 50
|
||||
blur_amount = 5
|
||||
gray = None
|
||||
|
||||
def sigmoid(x):
|
||||
return 1.0 / (1.0 + np.exp(-x))
|
||||
|
||||
def post_processing(output, shape):
|
||||
h, w = shape
|
||||
preds = []
|
||||
for p in output:
|
||||
img = sigmoid(p)
|
||||
img = np.squeeze(img)
|
||||
img = cv.normalize(img, None, 0, 255, cv.NORM_MINMAX, cv.CV_8U)
|
||||
img = cv.resize(img, (w, h))
|
||||
preds.append(img)
|
||||
fuse = preds[-1]
|
||||
ave = np.array(preds, dtype=np.float32)
|
||||
ave = np.uint8(np.mean(ave, axis=0))
|
||||
return fuse, ave
|
||||
|
||||
def apply_canny(image):
|
||||
global threshold1, threshold2, blur_amount
|
||||
kernel_size = 2 * blur_amount + 1
|
||||
blurred = cv.GaussianBlur(image, (kernel_size, kernel_size), 0)
|
||||
result = cv.Canny(blurred, threshold1, threshold2)
|
||||
cv.imshow('Output', result)
|
||||
|
||||
def setupCannyWindow(image):
|
||||
global gray
|
||||
cv.destroyWindow('Output')
|
||||
cv.namedWindow('Output', cv.WINDOW_AUTOSIZE)
|
||||
cv.moveWindow('Output', 200, 50)
|
||||
gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
|
||||
|
||||
cv.createTrackbar('thrs1', 'Output', threshold1, 255, lambda value: [globals().__setitem__('threshold1', value), apply_canny(gray)])
|
||||
cv.createTrackbar('thrs2', 'Output', threshold2, 255, lambda value: [globals().__setitem__('threshold2', value), apply_canny(gray)])
|
||||
cv.createTrackbar('blur', 'Output', blur_amount, 20, lambda value: [globals().__setitem__('blur_amount', value), apply_canny(gray)])
|
||||
|
||||
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 apply_dexined(model, image):
|
||||
t0 = cv.getTickCount()
|
||||
out = model.forward()
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
result,_ = post_processing(out, image.shape[:2])
|
||||
label = 'Inference time: %.2f ms' % (t * 1000.0)
|
||||
cv.putText(image, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
|
||||
cv.putText(result, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255))
|
||||
cv.imshow("Output", result)
|
||||
|
||||
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
|
||||
|
||||
cap = cv.VideoCapture(cv.samples.findFile(args.input) if args.input else 0)
|
||||
if not cap.isOpened():
|
||||
print("Failed to open the input video")
|
||||
exit(-1)
|
||||
cv.namedWindow('Input', cv.WINDOW_AUTOSIZE)
|
||||
cv.namedWindow('Output', cv.WINDOW_AUTOSIZE)
|
||||
cv.moveWindow('Output', 200, 50)
|
||||
|
||||
method = args.method
|
||||
if os.getenv('OPENCV_SAMPLES_DATA_PATH') is not None or hasattr(args, 'model'):
|
||||
try:
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
method = 'dexined'
|
||||
except:
|
||||
print("[WARN] Model file not provided, using canny instead. Pass model using --model=/path/to/dexined.onnx to use dexined model.")
|
||||
method = 'canny'
|
||||
args.model = None
|
||||
else:
|
||||
print("[WARN] Model file not provided, using canny instead. Pass model using --model=/path/to/dexined.onnx to use dexined model.")
|
||||
method = 'canny'
|
||||
|
||||
if method == 'canny':
|
||||
dummy = np.zeros((512, 512, 3), dtype="uint8")
|
||||
setupCannyWindow(dummy)
|
||||
net = None
|
||||
if method == "dexined":
|
||||
net = loadModel(args, engine)
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, image = cap.read()
|
||||
if not hasFrame:
|
||||
print("Press any key to exit")
|
||||
cv.waitKey(0)
|
||||
break
|
||||
if method == "canny":
|
||||
global gray
|
||||
gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
|
||||
apply_canny(gray)
|
||||
elif method == "dexined":
|
||||
inp = cv.dnn.blobFromImage(image, args.scale, (args.width, args.height), args.mean, swapRB=args.rgb, crop=False)
|
||||
|
||||
net.setInput(inp)
|
||||
apply_dexined(net, image)
|
||||
|
||||
cv.imshow("Input", image)
|
||||
key = cv.waitKey(30)
|
||||
if key == ord(' ') and method == 'canny':
|
||||
if hasattr(args, 'model') and args.model is not None:
|
||||
print("model: ", args.model)
|
||||
method = "dexined"
|
||||
if net is None:
|
||||
net = loadModel(args, engine)
|
||||
cv.destroyWindow('Output')
|
||||
cv.namedWindow('Output', cv.WINDOW_AUTOSIZE)
|
||||
cv.moveWindow('Output', 200, 50)
|
||||
else:
|
||||
print("[ERROR] Provide model file using --model to use dexined. Download model using python download_models.py dexined from dnn samples directory")
|
||||
elif key == ord(' ') and method=='dexined':
|
||||
method = "canny"
|
||||
setupCannyWindow(image)
|
||||
elif key == 27 or key == ord('q'):
|
||||
break
|
||||
cv.destroyAllWindows()
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,282 @@
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/objdetect.hpp>
|
||||
|
||||
#include <iostream>
|
||||
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
|
||||
static
|
||||
void visualize(Mat& input, int frame, Mat& faces, double fps, int thickness = 2)
|
||||
{
|
||||
std::string fpsString = cv::format("FPS : %.2f", (float)fps);
|
||||
if (frame >= 0)
|
||||
cout << "Frame " << frame << ", ";
|
||||
cout << "FPS: " << fpsString << endl;
|
||||
for (int i = 0; i < faces.rows; i++)
|
||||
{
|
||||
// Print results
|
||||
cout << "Face " << i
|
||||
<< ", top-left coordinates: (" << faces.at<float>(i, 0) << ", " << faces.at<float>(i, 1) << "), "
|
||||
<< "box width: " << faces.at<float>(i, 2) << ", box height: " << faces.at<float>(i, 3) << ", "
|
||||
<< "score: " << cv::format("%.2f", faces.at<float>(i, 14))
|
||||
<< endl;
|
||||
|
||||
// Draw bounding box
|
||||
rectangle(input, Rect2i(int(faces.at<float>(i, 0)), int(faces.at<float>(i, 1)), int(faces.at<float>(i, 2)), int(faces.at<float>(i, 3))), Scalar(0, 255, 0), thickness);
|
||||
// Draw landmarks
|
||||
circle(input, Point2i(int(faces.at<float>(i, 4)), int(faces.at<float>(i, 5))), 2, Scalar(255, 0, 0), thickness);
|
||||
circle(input, Point2i(int(faces.at<float>(i, 6)), int(faces.at<float>(i, 7))), 2, Scalar(0, 0, 255), thickness);
|
||||
circle(input, Point2i(int(faces.at<float>(i, 8)), int(faces.at<float>(i, 9))), 2, Scalar(0, 255, 0), thickness);
|
||||
circle(input, Point2i(int(faces.at<float>(i, 10)), int(faces.at<float>(i, 11))), 2, Scalar(255, 0, 255), thickness);
|
||||
circle(input, Point2i(int(faces.at<float>(i, 12)), int(faces.at<float>(i, 13))), 2, Scalar(0, 255, 255), thickness);
|
||||
}
|
||||
putText(input, fpsString, Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0), 2);
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
CommandLineParser parser(argc, argv,
|
||||
"{help h | | Print this message}"
|
||||
"{image1 i1 | | Path to the input image1. Omit for detecting through VideoCapture}"
|
||||
"{image2 i2 | | Path to the input image2. When image1 and image2 parameters given then the program try to find a face on both images and runs face recognition algorithm}"
|
||||
"{video v | 0 | Path to the input video}"
|
||||
"{scale sc | 1.0 | Scale factor used to resize input video frames}"
|
||||
"{fd_model fd | face_detection_yunet_2026may.onnx| Path to the model. Download yunet.onnx in https://github.com/opencv/opencv_zoo/tree/master/models/face_detection_yunet}"
|
||||
"{fr_model fr | face_recognition_sface_2021dec.onnx | Path to the face recognition model. Download the model at https://github.com/opencv/opencv_zoo/tree/master/models/face_recognition_sface}"
|
||||
"{score_threshold | 0.85 | Filter out faces of score < score_threshold}"
|
||||
"{nms_threshold | 0.3 | Suppress bounding boxes of iou >= nms_threshold}"
|
||||
"{top_k | 5000 | Keep top_k bounding boxes before NMS}"
|
||||
"{save s | false | Set true to save results. This flag is invalid when using camera}"
|
||||
);
|
||||
if (parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
String fd_modelPath = parser.get<String>("fd_model");
|
||||
String fr_modelPath = parser.get<String>("fr_model");
|
||||
|
||||
float scoreThreshold = parser.get<float>("score_threshold");
|
||||
float nmsThreshold = parser.get<float>("nms_threshold");
|
||||
int topK = parser.get<int>("top_k");
|
||||
|
||||
bool save = parser.get<bool>("save");
|
||||
float scale = parser.get<float>("scale");
|
||||
|
||||
double cosine_similar_thresh = 0.363;
|
||||
double l2norm_similar_thresh = 1.128;
|
||||
|
||||
//! [initialize_FaceDetectorYN]
|
||||
// Initialize FaceDetectorYN
|
||||
Ptr<FaceDetectorYN> detector = FaceDetectorYN::create(fd_modelPath, "", Size(320, 320), scoreThreshold, nmsThreshold, topK);
|
||||
//! [initialize_FaceDetectorYN]
|
||||
|
||||
TickMeter tm;
|
||||
|
||||
// If input is an image
|
||||
if (parser.has("image1"))
|
||||
{
|
||||
String input1 = parser.get<String>("image1");
|
||||
Mat image1 = imread(samples::findFile(input1));
|
||||
if (image1.empty())
|
||||
{
|
||||
std::cerr << "Cannot read image: " << input1 << std::endl;
|
||||
return 2;
|
||||
}
|
||||
|
||||
int imageWidth = int(image1.cols * scale);
|
||||
int imageHeight = int(image1.rows * scale);
|
||||
resize(image1, image1, Size(imageWidth, imageHeight));
|
||||
tm.start();
|
||||
|
||||
//! [inference]
|
||||
// Set input size before inference
|
||||
detector->setInputSize(image1.size());
|
||||
|
||||
Mat faces1;
|
||||
detector->detect(image1, faces1);
|
||||
if (faces1.rows < 1)
|
||||
{
|
||||
std::cerr << "Cannot find a face in " << input1 << std::endl;
|
||||
return 1;
|
||||
}
|
||||
//! [inference]
|
||||
|
||||
tm.stop();
|
||||
// Draw results on the input image
|
||||
visualize(image1, -1, faces1, tm.getFPS());
|
||||
|
||||
// Save results if save is true
|
||||
if (save)
|
||||
{
|
||||
cout << "Saving result.jpg...\n";
|
||||
imwrite("result.jpg", image1);
|
||||
}
|
||||
|
||||
// Visualize results
|
||||
imshow("image1", image1);
|
||||
pollKey(); // handle UI events to show content
|
||||
|
||||
if (parser.has("image2"))
|
||||
{
|
||||
String input2 = parser.get<String>("image2");
|
||||
Mat image2 = imread(samples::findFile(input2));
|
||||
if (image2.empty())
|
||||
{
|
||||
std::cerr << "Cannot read image2: " << input2 << std::endl;
|
||||
return 2;
|
||||
}
|
||||
|
||||
tm.reset();
|
||||
tm.start();
|
||||
detector->setInputSize(image2.size());
|
||||
|
||||
Mat faces2;
|
||||
detector->detect(image2, faces2);
|
||||
if (faces2.rows < 1)
|
||||
{
|
||||
std::cerr << "Cannot find a face in " << input2 << std::endl;
|
||||
return 1;
|
||||
}
|
||||
tm.stop();
|
||||
visualize(image2, -1, faces2, tm.getFPS());
|
||||
if (save)
|
||||
{
|
||||
cout << "Saving result2.jpg...\n";
|
||||
imwrite("result2.jpg", image2);
|
||||
}
|
||||
imshow("image2", image2);
|
||||
pollKey();
|
||||
|
||||
//! [initialize_FaceRecognizerSF]
|
||||
// Initialize FaceRecognizerSF
|
||||
Ptr<FaceRecognizerSF> faceRecognizer = FaceRecognizerSF::create(fr_modelPath, "");
|
||||
//! [initialize_FaceRecognizerSF]
|
||||
|
||||
|
||||
//! [facerecognizer]
|
||||
// Aligning and cropping facial image through the first face of faces detected.
|
||||
Mat aligned_face1, aligned_face2;
|
||||
faceRecognizer->alignCrop(image1, faces1.row(0), aligned_face1);
|
||||
faceRecognizer->alignCrop(image2, faces2.row(0), aligned_face2);
|
||||
|
||||
// Run feature extraction with given aligned_face
|
||||
Mat feature1, feature2;
|
||||
faceRecognizer->feature(aligned_face1, feature1);
|
||||
feature1 = feature1.clone();
|
||||
faceRecognizer->feature(aligned_face2, feature2);
|
||||
feature2 = feature2.clone();
|
||||
//! [facerecognizer]
|
||||
|
||||
//! [match]
|
||||
double cos_score = faceRecognizer->match(feature1, feature2, FaceRecognizerSF::DisType::FR_COSINE);
|
||||
double L2_score = faceRecognizer->match(feature1, feature2, FaceRecognizerSF::DisType::FR_NORM_L2);
|
||||
//! [match]
|
||||
|
||||
if (cos_score >= cosine_similar_thresh)
|
||||
{
|
||||
std::cout << "They have the same identity;";
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "They have different identities;";
|
||||
}
|
||||
std::cout << " Cosine Similarity: " << cos_score << ", threshold: " << cosine_similar_thresh << ". (higher value means higher similarity, max 1.0)\n";
|
||||
|
||||
if (L2_score <= l2norm_similar_thresh)
|
||||
{
|
||||
std::cout << "They have the same identity;";
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "They have different identities.";
|
||||
}
|
||||
std::cout << " NormL2 Distance: " << L2_score << ", threshold: " << l2norm_similar_thresh << ". (lower value means higher similarity, min 0.0)\n";
|
||||
}
|
||||
cout << "Press any key to exit..." << endl;
|
||||
waitKey(0);
|
||||
}
|
||||
else
|
||||
{
|
||||
int frameWidth, frameHeight;
|
||||
VideoCapture capture;
|
||||
std::string video = parser.get<string>("video");
|
||||
if (video.size() == 1 && isdigit(video[0]))
|
||||
capture.open(parser.get<int>("video"));
|
||||
else
|
||||
capture.open(samples::findFileOrKeep(video)); // keep GStreamer pipelines
|
||||
if (capture.isOpened())
|
||||
{
|
||||
frameWidth = int(capture.get(CAP_PROP_FRAME_WIDTH) * scale);
|
||||
frameHeight = int(capture.get(CAP_PROP_FRAME_HEIGHT) * scale);
|
||||
cout << "Video " << video
|
||||
<< ": width=" << frameWidth
|
||||
<< ", height=" << frameHeight
|
||||
<< endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
cout << "Could not initialize video capturing: " << video << "\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
detector->setInputSize(Size(frameWidth, frameHeight));
|
||||
|
||||
cout << "Press 'SPACE' to save frame, any other key to exit..." << endl;
|
||||
int nFrame = 0;
|
||||
for (;;)
|
||||
{
|
||||
// Get frame
|
||||
Mat frame;
|
||||
if (!capture.read(frame))
|
||||
{
|
||||
cerr << "Can't grab frame! Stop\n";
|
||||
break;
|
||||
}
|
||||
|
||||
resize(frame, frame, Size(frameWidth, frameHeight));
|
||||
|
||||
// Inference
|
||||
Mat faces;
|
||||
tm.start();
|
||||
detector->detect(frame, faces);
|
||||
tm.stop();
|
||||
|
||||
Mat result = frame.clone();
|
||||
// Draw results on the input image
|
||||
visualize(result, nFrame, faces, tm.getFPS());
|
||||
|
||||
// Visualize results
|
||||
imshow("Live", result);
|
||||
|
||||
int key = waitKey(1);
|
||||
bool saveFrame = save;
|
||||
if (key == ' ')
|
||||
{
|
||||
saveFrame = true;
|
||||
key = 0; // handled
|
||||
}
|
||||
|
||||
if (saveFrame)
|
||||
{
|
||||
std::string frame_name = cv::format("frame_%05d.png", nFrame);
|
||||
std::string result_name = cv::format("result_%05d.jpg", nFrame);
|
||||
cout << "Saving '" << frame_name << "' and '" << result_name << "' ...\n";
|
||||
imwrite(frame_name, frame);
|
||||
imwrite(result_name, result);
|
||||
}
|
||||
|
||||
++nFrame;
|
||||
|
||||
if (key > 0)
|
||||
break;
|
||||
}
|
||||
cout << "Processed " << nFrame << " frames" << endl;
|
||||
}
|
||||
cout << "Done." << endl;
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
|
||||
def str2bool(v):
|
||||
if v.lower() in ['on', 'yes', 'true', 'y', 't']:
|
||||
return True
|
||||
elif v.lower() in ['off', 'no', 'false', 'n', 'f']:
|
||||
return False
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--image1', '-i1', type=str, help='Path to the input image1. Omit for detecting on default camera.')
|
||||
parser.add_argument('--image2', '-i2', type=str, help='Path to the input image2. When image1 and image2 parameters given then the program try to find a face on both images and runs face recognition algorithm.')
|
||||
parser.add_argument('--video', '-v', type=str, help='Path to the input video.')
|
||||
parser.add_argument('--scale', '-sc', type=float, default=1.0, help='Scale factor used to resize input video frames.')
|
||||
parser.add_argument('--face_detection_model', '-fd', type=str, default='face_detection_yunet_2026may.onnx', help='Path to the face detection model. Download the model at https://github.com/opencv/opencv_zoo/tree/master/models/face_detection_yunet')
|
||||
parser.add_argument('--face_recognition_model', '-fr', type=str, default='face_recognition_sface_2021dec.onnx', help='Path to the face recognition model. Download the model at https://github.com/opencv/opencv_zoo/tree/master/models/face_recognition_sface')
|
||||
parser.add_argument('--score_threshold', type=float, default=0.85, help='Filtering out faces of score < score_threshold.')
|
||||
parser.add_argument('--nms_threshold', type=float, default=0.3, help='Suppress bounding boxes of iou >= nms_threshold.')
|
||||
parser.add_argument('--top_k', type=int, default=5000, help='Keep top_k bounding boxes before NMS.')
|
||||
parser.add_argument('--save', '-s', type=str2bool, default=False, help='Set true to save results. This flag is invalid when using camera.')
|
||||
args = parser.parse_args()
|
||||
|
||||
def visualize(input, faces, fps, thickness=2):
|
||||
if faces[1] is not None:
|
||||
for idx, face in enumerate(faces[1]):
|
||||
print('Face {}, top-left coordinates: ({:.0f}, {:.0f}), box width: {:.0f}, box height {:.0f}, score: {:.2f}'.format(idx, face[0], face[1], face[2], face[3], face[-1]))
|
||||
|
||||
coords = face[:-1].astype(np.int32)
|
||||
cv.rectangle(input, (coords[0], coords[1]), (coords[0]+coords[2], coords[1]+coords[3]), (0, 255, 0), thickness)
|
||||
cv.circle(input, (coords[4], coords[5]), 2, (255, 0, 0), thickness)
|
||||
cv.circle(input, (coords[6], coords[7]), 2, (0, 0, 255), thickness)
|
||||
cv.circle(input, (coords[8], coords[9]), 2, (0, 255, 0), thickness)
|
||||
cv.circle(input, (coords[10], coords[11]), 2, (255, 0, 255), thickness)
|
||||
cv.circle(input, (coords[12], coords[13]), 2, (0, 255, 255), thickness)
|
||||
cv.putText(input, 'FPS: {:.2f}'.format(fps), (1, 16), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
## [initialize_FaceDetectorYN]
|
||||
detector = cv.FaceDetectorYN.create(
|
||||
args.face_detection_model,
|
||||
"",
|
||||
(320, 320),
|
||||
args.score_threshold,
|
||||
args.nms_threshold,
|
||||
args.top_k
|
||||
)
|
||||
## [initialize_FaceDetectorYN]
|
||||
|
||||
tm = cv.TickMeter()
|
||||
|
||||
# If input is an image
|
||||
if args.image1 is not None:
|
||||
img1 = cv.imread(cv.samples.findFile(args.image1))
|
||||
img1Width = int(img1.shape[1]*args.scale)
|
||||
img1Height = int(img1.shape[0]*args.scale)
|
||||
|
||||
img1 = cv.resize(img1, (img1Width, img1Height))
|
||||
tm.start()
|
||||
|
||||
## [inference]
|
||||
# Set input size before inference
|
||||
detector.setInputSize((img1Width, img1Height))
|
||||
|
||||
faces1 = detector.detect(img1)
|
||||
## [inference]
|
||||
|
||||
tm.stop()
|
||||
assert faces1[1] is not None, 'Cannot find a face in {}'.format(args.image1)
|
||||
|
||||
# Draw results on the input image
|
||||
visualize(img1, faces1, tm.getFPS())
|
||||
|
||||
# Save results if save is true
|
||||
if args.save:
|
||||
print('Results saved to result.jpg\n')
|
||||
cv.imwrite('result.jpg', img1)
|
||||
|
||||
# Visualize results in a new window
|
||||
cv.imshow("image1", img1)
|
||||
|
||||
if args.image2 is not None:
|
||||
img2 = cv.imread(cv.samples.findFile(args.image2))
|
||||
|
||||
tm.reset()
|
||||
tm.start()
|
||||
detector.setInputSize((img2.shape[1], img2.shape[0]))
|
||||
faces2 = detector.detect(img2)
|
||||
tm.stop()
|
||||
assert faces2[1] is not None, 'Cannot find a face in {}'.format(args.image2)
|
||||
visualize(img2, faces2, tm.getFPS())
|
||||
cv.imshow("image2", img2)
|
||||
|
||||
## [initialize_FaceRecognizerSF]
|
||||
recognizer = cv.FaceRecognizerSF.create(
|
||||
args.face_recognition_model,"")
|
||||
## [initialize_FaceRecognizerSF]
|
||||
|
||||
## [facerecognizer]
|
||||
# Align faces
|
||||
face1_align = recognizer.alignCrop(img1, faces1[1][0])
|
||||
face2_align = recognizer.alignCrop(img2, faces2[1][0])
|
||||
|
||||
# Extract features
|
||||
face1_feature = recognizer.feature(face1_align)
|
||||
face2_feature = recognizer.feature(face2_align)
|
||||
## [facerecognizer]
|
||||
|
||||
cosine_similarity_threshold = 0.363
|
||||
l2_similarity_threshold = 1.128
|
||||
|
||||
## [match]
|
||||
cosine_score = recognizer.match(face1_feature, face2_feature, cv.FaceRecognizerSF_FR_COSINE)
|
||||
l2_score = recognizer.match(face1_feature, face2_feature, cv.FaceRecognizerSF_FR_NORM_L2)
|
||||
## [match]
|
||||
|
||||
msg = 'different identities'
|
||||
if cosine_score >= cosine_similarity_threshold:
|
||||
msg = 'the same identity'
|
||||
print('They have {}. Cosine Similarity: {}, threshold: {} (higher value means higher similarity, max 1.0).'.format(msg, cosine_score, cosine_similarity_threshold))
|
||||
|
||||
msg = 'different identities'
|
||||
if l2_score <= l2_similarity_threshold:
|
||||
msg = 'the same identity'
|
||||
print('They have {}. NormL2 Distance: {}, threshold: {} (lower value means higher similarity, min 0.0).'.format(msg, l2_score, l2_similarity_threshold))
|
||||
cv.waitKey(0)
|
||||
else: # Omit input to call default camera
|
||||
if args.video is not None:
|
||||
deviceId = args.video
|
||||
else:
|
||||
deviceId = 0
|
||||
cap = cv.VideoCapture(deviceId)
|
||||
frameWidth = int(cap.get(cv.CAP_PROP_FRAME_WIDTH)*args.scale)
|
||||
frameHeight = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)*args.scale)
|
||||
detector.setInputSize([frameWidth, frameHeight])
|
||||
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
print('No frames grabbed!')
|
||||
break
|
||||
|
||||
frame = cv.resize(frame, (frameWidth, frameHeight))
|
||||
|
||||
# Inference
|
||||
tm.start()
|
||||
faces = detector.detect(frame) # faces is a tuple
|
||||
tm.stop()
|
||||
|
||||
# Draw results on the input image
|
||||
visualize(frame, faces, tm.getFPS())
|
||||
|
||||
# Visualize results
|
||||
cv.imshow('Live', frame)
|
||||
cv.destroyAllWindows()
|
||||
@@ -0,0 +1,2 @@
|
||||
*.caffemodel
|
||||
*.pb
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,79 @@
|
||||
This is a brief description of training process which has been used to get res10_300x300_ssd_iter_140000.caffemodel.
|
||||
The model was created with SSD framework using ResNet-10 like architecture as a backbone. Channels count in ResNet-10 convolution layers was significantly dropped (2x- or 4x- fewer channels).
|
||||
The model was trained in Caffe framework on some huge and available online dataset.
|
||||
|
||||
1. Prepare training tools
|
||||
You need to use "ssd" branch from this repository https://github.com/weiliu89/caffe/tree/ssd . Checkout this branch and built it (see instructions in repo's README)
|
||||
|
||||
2. Prepare training data.
|
||||
The data preparation pipeline can be represented as:
|
||||
|
||||
(a)Download original face detection dataset -> (b)Convert annotation to the PASCAL VOC format -> (c)Create LMDB database with images + annotations for training
|
||||
|
||||
a) Find some datasets with face bounding boxes annotation. For some reasons I can't provide links here, but you easily find them on your own. Also study the data. It may contain small or low quality faces which can spoil training process. Often there are special flags about object quality in annotation. Remove such faces from annotation (smaller when 16 along at least one side, or blurred, of highly-occluded, or something else).
|
||||
|
||||
b) The downloaded dataset will have some format of annotation. It may be one single file for all images, or separate file for each image or something else. But to train SSD in Caffe you need to convert annotation to PASCAL VOC format.
|
||||
PASCAL VOC annotation consist of .xml file for each image. In this xml file all face bounding boxes should be listed as:
|
||||
|
||||
<annotation>
|
||||
<size>
|
||||
<width>300</width>
|
||||
<height>300</height>
|
||||
</size>
|
||||
<object>
|
||||
<name>face</name>
|
||||
<difficult>0</difficult>
|
||||
<bndbox>
|
||||
<xmin>100</xmin>
|
||||
<ymin>100</ymin>
|
||||
<xmax>200</xmax>
|
||||
<ymax>200</ymax>
|
||||
</bndbox>
|
||||
</object>
|
||||
<object>
|
||||
<name>face</name>
|
||||
<difficult>0</difficult>
|
||||
<bndbox>
|
||||
<xmin>0</xmin>
|
||||
<ymin>0</ymin>
|
||||
<xmax>100</xmax>
|
||||
<ymax>100</ymax>
|
||||
</bndbox>
|
||||
</object>
|
||||
</annotation>
|
||||
|
||||
So, convert your dataset's annotation to the format above.
|
||||
Also, you should create labelmap.prototxt file with the following content:
|
||||
item {
|
||||
name: "none_of_the_above"
|
||||
label: 0
|
||||
display_name: "background"
|
||||
}
|
||||
item {
|
||||
name: "face"
|
||||
label: 1
|
||||
display_name: "face"
|
||||
}
|
||||
|
||||
You need this file to establish correspondence between name of class and digital label of class.
|
||||
|
||||
For next step we also need file there all our image-annotation file names pairs are listed. This file should contain similar lines:
|
||||
images_val/0.jpg annotations_val/0.jpg.xml
|
||||
|
||||
c) To create LMDB you need to use create_data.sh tool from caffe/data/VOC0712 Caffe's source code directory.
|
||||
This script calls create_annoset.py inside, so check out what you need to pass as script's arguments
|
||||
|
||||
You need to prepare 2 LMDB databases: one for training images, one for validation images.
|
||||
|
||||
3. Train your detector
|
||||
For training you need to have 3 files: train.prototxt, test.prototxt and solver.prototxt. You can find these files in the same directory as for this readme.
|
||||
Also you need to edit train.prototxt and test.prototxt to replace paths for your LMDB databases to actual databases you've created in step 2.
|
||||
|
||||
Now all is done for launch training process.
|
||||
Execute next lines in Terminal:
|
||||
mkdir -p snapshot
|
||||
mkdir -p log
|
||||
/path_for_caffe_build_dir/tools/caffe train -solver="solver.prototxt" -gpu 0 2>&1 | tee -a log/log.log
|
||||
|
||||
And wait. It will take about 8 hours to finish the process.
|
||||
After it you can use your .caffemodel from snapshot/ subdirectory in resnet_face_ssd_python.py sample.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,28 @@
|
||||
train_net: "train.prototxt"
|
||||
test_net: "test.prototxt"
|
||||
|
||||
test_iter: 2312
|
||||
test_interval: 5000
|
||||
test_initialization: true
|
||||
|
||||
base_lr: 0.01
|
||||
display: 10
|
||||
lr_policy: "multistep"
|
||||
max_iter: 140000
|
||||
stepvalue: 80000
|
||||
stepvalue: 120000
|
||||
gamma: 0.1
|
||||
momentum: 0.9
|
||||
weight_decay: 0.0005
|
||||
average_loss: 500
|
||||
iter_size: 1
|
||||
type: "SGD"
|
||||
|
||||
solver_mode: GPU
|
||||
random_seed: 0
|
||||
debug_info: false
|
||||
snapshot: 1000
|
||||
snapshot_prefix: "snapshot/res10_300x300_ssd"
|
||||
|
||||
eval_type: "detection"
|
||||
ap_version: "11point"
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,13 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<metalink xmlns="urn:ietf:params:xml:ns:metalink">
|
||||
<file name="res10_300x300_ssd_iter_140000_fp16.caffemodel">
|
||||
<identity>opencv_face_detector_fp16</identity>
|
||||
<hash type="sha-1">31fc22bfdd907567a04bb45b7cfad29966caddc1</hash>
|
||||
<url>https://raw.githubusercontent.com/opencv/opencv_3rdparty/dnn_samples_face_detector_20180205_fp16/res10_300x300_ssd_iter_140000_fp16.caffemodel</url>
|
||||
</file>
|
||||
<file name="opencv_face_detector_uint8.pb">
|
||||
<identity>opencv_face_detector_uint8</identity>
|
||||
<hash type="sha-1">4f2fdf6f231d759d7bbdb94353c5a68690f3d2ae</hash>
|
||||
<url>https://raw.githubusercontent.com/opencv/opencv_3rdparty/dnn_samples_face_detector_20180220_uint8/opencv_face_detector_uint8.pb</url>
|
||||
</file>
|
||||
</metalink>
|
||||
@@ -0,0 +1,52 @@
|
||||
from __future__ import print_function
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description='This script is used to run style transfer models from '
|
||||
'https://github.com/onnx/models/tree/main/vision/style_transfer/fast_neural_style using OpenCV')
|
||||
parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera')
|
||||
parser.add_argument('--model', help='Path to .onnx model')
|
||||
parser.add_argument('--width', default=-1, type=int, help='Resize input to specific width.')
|
||||
parser.add_argument('--height', default=-1, type=int, help='Resize input to specific height.')
|
||||
parser.add_argument('--median_filter', default=0, type=int, help='Kernel size of postprocessing blurring.')
|
||||
args = parser.parse_args()
|
||||
|
||||
net = cv.dnn.readNet(cv.samples.findFile(args.model))
|
||||
net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
|
||||
|
||||
if args.input:
|
||||
cap = cv.VideoCapture(args.input)
|
||||
else:
|
||||
cap = cv.VideoCapture(0)
|
||||
|
||||
cv.namedWindow('Styled image', cv.WINDOW_NORMAL)
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
cv.waitKey()
|
||||
break
|
||||
|
||||
inWidth = args.width if args.width != -1 else frame.shape[1]
|
||||
inHeight = args.height if args.height != -1 else frame.shape[0]
|
||||
inp = cv.dnn.blobFromImage(frame, 1.0, (inWidth, inHeight),
|
||||
swapRB=True, crop=False)
|
||||
|
||||
net.setInput(inp)
|
||||
t0 = cv.getTickCount()
|
||||
out = net.forward()
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
|
||||
out = out.reshape(3, out.shape[2], out.shape[3])
|
||||
out = out.transpose(1, 2, 0)
|
||||
|
||||
print('%.2f ms' % (t * 1000.0))
|
||||
|
||||
if args.median_filter:
|
||||
out = cv.medianBlur(out, args.median_filter)
|
||||
|
||||
out = np.clip(out, 0, 255)
|
||||
out = out.astype(np.uint8)
|
||||
|
||||
cv.imshow('Styled image', out)
|
||||
@@ -0,0 +1,140 @@
|
||||
# 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) 2026, BigVision LLC, all rights reserved.
|
||||
# Third party copyrights are property of their respective owners.
|
||||
|
||||
'''
|
||||
This is a sample script to run Gemma3 inference in OpenCV using ONNX model.
|
||||
The script loads the Gemma3 model and runs inference on a given prompt using
|
||||
the Gemma3 chat format (<start_of_turn> / <end_of_turn> special tokens).
|
||||
|
||||
Model: https://huggingface.co/google/gemma-3-1b-it
|
||||
|
||||
Exporting Gemma3 model to ONNX:
|
||||
|
||||
1. Install the required dependencies:
|
||||
|
||||
pip install optimum[exporters] optimum-onnx[onnxruntime] torch transformers
|
||||
|
||||
2. Export the model to ONNX:
|
||||
|
||||
Without KV-cache:
|
||||
|
||||
optimum-cli export onnx --model google/gemma-3-1b-it --task causal-lm gemma3_instruct_onnx/
|
||||
|
||||
With KV-cache (recommended, faster autoregressive inference):
|
||||
|
||||
optimum-cli export onnx --model google/gemma-3-1b-it --task causal-lm-with-past gemma3_instruct_onnx_with_past/
|
||||
|
||||
|
||||
Run the script:
|
||||
1. Install the required dependencies:
|
||||
|
||||
pip install numpy
|
||||
|
||||
2. Run the script:
|
||||
|
||||
Without KV-cache (causal-lm export):
|
||||
|
||||
python gemma3_inference.py --model=<path-to-onnx-model> \
|
||||
--tokenizer_path=<path-to-opencv-tokenizer-config.json> \
|
||||
--prompt="What is OpenCV?"
|
||||
|
||||
With KV-cache (causal-lm-with-past export):
|
||||
|
||||
python gemma3_inference.py --model=<path-to-onnx-model> \
|
||||
--tokenizer_path=<path-to-opencv-tokenizer-config.json> \
|
||||
--prompt="What is OpenCV?" \
|
||||
--use_kv_cache
|
||||
|
||||
The tokenizer_path should point to an OpenCV-format config.json (e.g., from
|
||||
opencv_extra/testdata/dnn/llm/gemma3/config.json), NOT the HuggingFace tokenizer_config.json.
|
||||
'''
|
||||
|
||||
import numpy as np
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Use this script to run Gemma3 inference in OpenCV',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--model', type=str, required=True, help='Path to Gemma3 ONNX model file.')
|
||||
parser.add_argument('--tokenizer_path', type=str, required=True, help='Path to Gemma3 tokenizer config.json.')
|
||||
parser.add_argument('--prompt', type=str, default='What is OpenCV?', help='User prompt.')
|
||||
parser.add_argument('--max_new_tokens', type=int, default=64, help='Maximum number of new tokens to generate.')
|
||||
parser.add_argument('--use_kv_cache', action='store_true', default=False, help='Enable KV-cache for faster inference (requires causal-lm-with-past export).')
|
||||
parser.add_argument('--seed', type=int, default=0, help='Random seed.')
|
||||
return parser.parse_args()
|
||||
|
||||
def build_gemma3_prompt(user_prompt):
|
||||
'''Wrap user prompt in Gemma3 chat format.'''
|
||||
return '<start_of_turn>user\n' + user_prompt + '<end_of_turn>\n<start_of_turn>model\n'
|
||||
|
||||
def gemma3_inference(net, prompt, max_new_tokens, tokenizer, use_kv_cache=True):
|
||||
|
||||
print("Inferencing Gemma3 model...")
|
||||
|
||||
tokens = tokenizer.encode(prompt)
|
||||
# Prepend BOS token (id=2) as required by Gemma3
|
||||
tokens = [2] + list(tokens)
|
||||
input_ids = np.array(tokens, dtype=np.int64).reshape(1, -1)
|
||||
|
||||
# Gemma3 special token IDs
|
||||
eos_id = 1 # <eos>
|
||||
eot_id = 106 # <end_of_turn>
|
||||
stop_ids = (eos_id, eot_id)
|
||||
|
||||
generated = []
|
||||
|
||||
if use_kv_cache:
|
||||
net.enableKVCache()
|
||||
prompt_len = input_ids.shape[1]
|
||||
|
||||
# Prefill: process full prompt once to populate KV-cache
|
||||
net.setInput(input_ids, 'input_ids')
|
||||
net.setInput(np.ones((1, prompt_len), dtype=np.int64), 'attention_mask')
|
||||
logits = net.forward()
|
||||
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
|
||||
generated = [new_id]
|
||||
|
||||
# Generate: feed one new token per step; OpenCV routes present.* -> past_key_values.*
|
||||
for _ in range(max_new_tokens - 1):
|
||||
if new_id in stop_ids:
|
||||
break
|
||||
net.setInput(np.array([[new_id]], dtype=np.int64), 'input_ids')
|
||||
net.setInput(np.ones((1, prompt_len + len(generated)), dtype=np.int64), 'attention_mask')
|
||||
logits = net.forward()
|
||||
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
|
||||
generated.append(new_id)
|
||||
else:
|
||||
# Without KV-cache: feed full growing sequence each step
|
||||
for _ in range(max_new_tokens):
|
||||
net.setInput(input_ids, 'input_ids')
|
||||
net.setInput(np.ones((1, input_ids.shape[1]), dtype=np.int64), 'attention_mask')
|
||||
logits = net.forward()
|
||||
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
|
||||
if new_id in stop_ids:
|
||||
break
|
||||
generated.append(new_id)
|
||||
input_ids = np.concatenate([input_ids, [[new_id]]], axis=1)
|
||||
|
||||
return np.array([tokens + generated], dtype=np.int64)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parse_args()
|
||||
np.random.seed(args.seed)
|
||||
|
||||
print("Preparing Gemma3 model...")
|
||||
tokenizer = cv.dnn.Tokenizer.load(args.tokenizer_path)
|
||||
|
||||
net = cv.dnn.readNetFromONNX(args.model, cv.dnn.ENGINE_NEW)
|
||||
|
||||
gemma3_prompt = build_gemma3_prompt(args.prompt)
|
||||
print(f"Prompt:\n{gemma3_prompt}")
|
||||
|
||||
prompt_len = len(tokenizer.encode(gemma3_prompt)) + 1 # +1 for BOS token
|
||||
tokens = gemma3_inference(net, gemma3_prompt, args.max_new_tokens, tokenizer, args.use_kv_cache)
|
||||
response = tokenizer.decode(tokens[0][prompt_len:].tolist())
|
||||
print(f"Response:\n{response}")
|
||||
@@ -0,0 +1,89 @@
|
||||
'''
|
||||
This is a sample script to run GPT-2 inference in OpenCV using ONNX model.
|
||||
The script loads the GPT-2 model and runs inference on a given prompt.
|
||||
Currently script only works with fixed size window, that means
|
||||
you will have to specify prompt of the same length as when model was exported to ONNX.
|
||||
|
||||
|
||||
Exporting GPT-2 model to ONNX.
|
||||
To export GPT-2 model to ONNX, you can use the following procedure:
|
||||
|
||||
1. Clone fork of Andrej Karpathy's GPT-2 repository:
|
||||
|
||||
git clone -b fix-dynamic-axis-export https://github.com/nklskyoy/build-nanogpt
|
||||
|
||||
2. Install the required dependencies:
|
||||
|
||||
pip install -r requirements.txt
|
||||
|
||||
3 Export the model to ONNX:
|
||||
|
||||
python export2onnx.py --promt=<Any-promt-you-want>
|
||||
|
||||
|
||||
Run the script:
|
||||
1. Install the required dependencies:
|
||||
|
||||
pip install tiktoken==0.7.0 numpy tqdm
|
||||
|
||||
2. Run the script:
|
||||
python gpt2_inference.py --model=<path-to-onnx-model> --tokenizer_path=<path-to-tokenizer-config> --prompt=<use-promt-of-the-same-length-used-while-exporting>
|
||||
'''
|
||||
|
||||
import numpy as np
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Use this script to run GPT-2 inference in OpenCV',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--model', type=str, required=True, help='Path to GPT-2 model ONNX model file.')
|
||||
parser.add_argument('--tokenizer_path', type=str, required=True, help='Path to GPT-2 tokenizer config file.')
|
||||
parser.add_argument("--prompt", type=str, default="Hello, I'm a language model,", help="Prompt to start with.")
|
||||
parser.add_argument("--max_seq_len", type=int, default=32, help="Number of tokens to continue.")
|
||||
parser.add_argument("--seed", type=int, default=0, help="Random seed")
|
||||
return parser.parse_args()
|
||||
|
||||
def stable_softmax(logits):
|
||||
exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
|
||||
return exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
|
||||
|
||||
|
||||
|
||||
def gpt2_inference(net, prompt, max_length, tokenizer):
|
||||
|
||||
print("Inferencing GPT-2 model...")
|
||||
|
||||
tokens = tokenizer.encode(prompt).reshape(1,-1)
|
||||
|
||||
stop_tokens = (50256, ) ## could be extended to include more stop tokens
|
||||
while 0 < max_length and tokens[:, -1] not in stop_tokens:
|
||||
|
||||
net.setInputsNames(['idx'])
|
||||
net.setInput(tokens, 'idx')
|
||||
logits = net.forward()
|
||||
logits = logits[:, -1, :] # (B, vocab_size)
|
||||
|
||||
# use hard sampling
|
||||
new_idx = np.argmax(logits.reshape(-1)).reshape(1,1)
|
||||
|
||||
tokens = np.concatenate((tokens, new_idx), axis=1)
|
||||
|
||||
max_length -= 1
|
||||
return tokens
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parse_args()
|
||||
print("Preparing GPT-2 model...")
|
||||
max_length = args.max_seq_len
|
||||
prompt = args.prompt
|
||||
tokenizer_path = args.tokenizer_path
|
||||
|
||||
net = cv.dnn.readNetFromONNX(args.model, cv.dnn.ENGINE_NEW)
|
||||
tokenizer = cv.dnn.Tokenizer.load(tokenizer_path)
|
||||
|
||||
tokens = gpt2_inference(net, prompt, max_length, tokenizer)
|
||||
print(tokenizer.decode(tokens[0]))
|
||||
@@ -0,0 +1,111 @@
|
||||
//
|
||||
// this sample demonstrates parsing (segmenting) human body parts from an image using opencv's dnn,
|
||||
// based on https://github.com/Engineering-Course/LIP_JPPNet
|
||||
//
|
||||
// get the pretrained model from: https://www.dropbox.com/s/qag9vzambhhkvxr/lip_jppnet_384.pb?dl=0
|
||||
//
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
using namespace cv;
|
||||
|
||||
|
||||
static Mat parse_human(const Mat &image, const std::string &model, int backend=dnn::DNN_BACKEND_DEFAULT, int target=dnn::DNN_TARGET_CPU) {
|
||||
// this network expects an image and a flipped copy as input
|
||||
Mat flipped;
|
||||
flip(image, flipped, 1);
|
||||
std::vector<Mat> batch;
|
||||
batch.push_back(image);
|
||||
batch.push_back(flipped);
|
||||
Mat blob = dnn::blobFromImages(batch, 1.0, Size(), Scalar(104.00698793, 116.66876762, 122.67891434));
|
||||
|
||||
dnn::Net net = dnn::readNet(model);
|
||||
net.setPreferableBackend(backend);
|
||||
net.setPreferableTarget(target);
|
||||
net.setInput(blob);
|
||||
Mat out = net.forward();
|
||||
// expected output: [2, 20, 384, 384], (2 lists(orig, flipped) of 20 body part heatmaps 384x384)
|
||||
|
||||
// LIP classes:
|
||||
// 0 Background, 1 Hat, 2 Hair, 3 Glove, 4 Sunglasses, 5 UpperClothes, 6 Dress, 7 Coat, 8 Socks, 9 Pants
|
||||
// 10 Jumpsuits, 11 Scarf, 12 Skirt, 13 Face, 14 LeftArm, 15 RightArm, 16 LeftLeg, 17 RightLeg, 18 LeftShoe. 19 RightShoe
|
||||
Vec3b colors[] = {
|
||||
Vec3b(0, 0, 0), Vec3b(128, 0, 0), Vec3b(255, 0, 0), Vec3b(0, 85, 0), Vec3b(170, 0, 51), Vec3b(255, 85, 0),
|
||||
Vec3b(0, 0, 85), Vec3b(0, 119, 221), Vec3b(85, 85, 0), Vec3b(0, 85, 85), Vec3b(85, 51, 0), Vec3b(52, 86, 128),
|
||||
Vec3b(0, 128, 0), Vec3b(0, 0, 255), Vec3b(51, 170, 221), Vec3b(0, 255, 255), Vec3b(85, 255, 170),
|
||||
Vec3b(170, 255, 85), Vec3b(255, 255, 0), Vec3b(255, 170, 0)
|
||||
};
|
||||
|
||||
Mat segm(image.size(), CV_8UC3, Scalar(0,0,0));
|
||||
Mat maxval(image.size(), CV_32F, Scalar(0));
|
||||
|
||||
// iterate over body part heatmaps (LIP classes)
|
||||
for (int i=0; i<out.size[1]; i++) {
|
||||
// resize heatmaps to original image size
|
||||
// "head" is the original image result, "tail" the flipped copy
|
||||
Mat head, h(out.size[2], out.size[3], CV_32F, out.ptr<float>(0,i));
|
||||
resize(h, head, image.size());
|
||||
|
||||
// we have to swap the last 3 pairs in the "tail" list
|
||||
static int tail_order[] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,15,14,17,16,19,18};
|
||||
Mat tail, t(out.size[2], out.size[3], CV_32F, out.ptr<float>(1,tail_order[i]));
|
||||
resize(t, tail, image.size());
|
||||
flip(tail, tail, 1);
|
||||
|
||||
// mix original and flipped result
|
||||
Mat avg = (head + tail) * 0.5;
|
||||
|
||||
// write color if prob value > maxval
|
||||
Mat cmask;
|
||||
compare(avg, maxval, cmask, CMP_GT);
|
||||
segm.setTo(colors[i], cmask);
|
||||
|
||||
// keep largest values for next iteration
|
||||
max(avg, maxval, maxval);
|
||||
}
|
||||
cvtColor(segm, segm, COLOR_RGB2BGR);
|
||||
return segm;
|
||||
}
|
||||
|
||||
int main(int argc, char**argv)
|
||||
{
|
||||
std::string param_keys =
|
||||
"{help h | | show help screen / args}"
|
||||
"{image i | | person image to process }"
|
||||
"{model m |lip_jppnet_384.pb| network model}";
|
||||
std::string backend_keys = cv::format(
|
||||
"{ backend | 0 | Choose one of computation backends: "
|
||||
"%d: automatically (by default), "
|
||||
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"%d: OpenCV implementation, "
|
||||
"%d: VKCOM, "
|
||||
"%d: CUDA }", cv::dnn::DNN_BACKEND_DEFAULT, cv::dnn::DNN_BACKEND_INFERENCE_ENGINE, cv::dnn::DNN_BACKEND_OPENCV, cv::dnn::DNN_BACKEND_VKCOM, cv::dnn::DNN_BACKEND_CUDA);
|
||||
std::string target_keys = cv::format(
|
||||
"{ target | 0 | Choose one of target computation devices: "
|
||||
"%d: CPU target (by default), "
|
||||
"%d: OpenCL, "
|
||||
"%d: OpenCL fp16 (half-float precision), "
|
||||
"%d: VPU, "
|
||||
"%d: Vulkan, "
|
||||
"%d: CUDA, "
|
||||
"%d: CUDA fp16 (half-float preprocess) }", cv::dnn::DNN_TARGET_CPU, cv::dnn::DNN_TARGET_OPENCL, cv::dnn::DNN_TARGET_OPENCL_FP16, cv::dnn::DNN_TARGET_MYRIAD, cv::dnn::DNN_TARGET_VULKAN, cv::dnn::DNN_TARGET_CUDA, cv::dnn::DNN_TARGET_CUDA_FP16);
|
||||
std::string keys = param_keys + backend_keys + target_keys;
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
if (argc == 1 || parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
std::string model = parser.get<std::string>("model");
|
||||
std::string image = parser.get<std::string>("image");
|
||||
int backend = parser.get<int>("backend");
|
||||
int target = parser.get<int>("target");
|
||||
|
||||
Mat input = imread(image);
|
||||
Mat segm = parse_human(input, model, backend, target);
|
||||
|
||||
imshow("human parsing", segm);
|
||||
waitKey();
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
#!/usr/bin/env python
|
||||
'''
|
||||
You can download the converted pb model from https://www.dropbox.com/s/qag9vzambhhkvxr/lip_jppnet_384.pb?dl=0
|
||||
or convert the model yourself.
|
||||
|
||||
Follow these steps if you want to convert the original model yourself:
|
||||
To get original .meta pre-trained model download https://drive.google.com/file/d/1BFVXgeln-bek8TCbRjN6utPAgRE0LJZg/view
|
||||
For correct convert .meta to .pb model download original repository https://github.com/Engineering-Course/LIP_JPPNet
|
||||
Change script evaluate_parsing_JPPNet-s2.py for human parsing
|
||||
1. Remove preprocessing to create image_batch_origin:
|
||||
with tf.name_scope("create_inputs"):
|
||||
...
|
||||
Add
|
||||
image_batch_origin = tf.placeholder(tf.float32, shape=(2, None, None, 3), name='input')
|
||||
|
||||
2. Create input
|
||||
image = cv2.imread(path/to/image)
|
||||
image_rev = np.flip(image, axis=1)
|
||||
input = np.stack([image, image_rev], axis=0)
|
||||
|
||||
3. Hardcode image_h and image_w shapes to determine output shapes.
|
||||
We use default INPUT_SIZE = (384, 384) from evaluate_parsing_JPPNet-s2.py.
|
||||
parsing_out1 = tf.reduce_mean(tf.stack([tf.image.resize_images(parsing_out1_100, INPUT_SIZE),
|
||||
tf.image.resize_images(parsing_out1_075, INPUT_SIZE),
|
||||
tf.image.resize_images(parsing_out1_125, INPUT_SIZE)]), axis=0)
|
||||
Do similarly with parsing_out2, parsing_out3
|
||||
4. Remove postprocessing. Last net operation:
|
||||
raw_output = tf.reduce_mean(tf.stack([parsing_out1, parsing_out2, parsing_out3]), axis=0)
|
||||
Change:
|
||||
parsing_ = sess.run(raw_output, feed_dict={'input:0': input})
|
||||
|
||||
5. To save model after sess.run(...) add:
|
||||
input_graph_def = tf.get_default_graph().as_graph_def()
|
||||
output_node = "Mean_3"
|
||||
output_graph_def = tf.graph_util.convert_variables_to_constants(sess, input_graph_def, output_node)
|
||||
|
||||
output_graph = "LIP_JPPNet.pb"
|
||||
with tf.gfile.GFile(output_graph, "wb") as f:
|
||||
f.write(output_graph_def.SerializeToString())'
|
||||
'''
|
||||
|
||||
import argparse
|
||||
import os.path
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
|
||||
|
||||
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV,
|
||||
cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
|
||||
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD,
|
||||
cv.dnn.DNN_TARGET_HDDL, cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)
|
||||
|
||||
|
||||
def preprocess(image):
|
||||
"""
|
||||
Create 4-dimensional blob from image and flip image
|
||||
:param image: input image
|
||||
"""
|
||||
image_rev = np.flip(image, axis=1)
|
||||
input = cv.dnn.blobFromImages([image, image_rev], mean=(104.00698793, 116.66876762, 122.67891434))
|
||||
return input
|
||||
|
||||
|
||||
def run_net(input, model_path, backend, target):
|
||||
"""
|
||||
Read network and infer model
|
||||
:param model_path: path to JPPNet model
|
||||
:param backend: computation backend
|
||||
:param target: computation device
|
||||
"""
|
||||
net = cv.dnn.readNet(model_path)
|
||||
net.setPreferableBackend(backend)
|
||||
net.setPreferableTarget(target)
|
||||
net.setInput(input)
|
||||
out = net.forward()
|
||||
return out
|
||||
|
||||
|
||||
def postprocess(out, input_shape):
|
||||
"""
|
||||
Create a grayscale human segmentation
|
||||
:param out: network output
|
||||
:param input_shape: input image width and height
|
||||
"""
|
||||
# LIP classes
|
||||
# 0 Background
|
||||
# 1 Hat
|
||||
# 2 Hair
|
||||
# 3 Glove
|
||||
# 4 Sunglasses
|
||||
# 5 UpperClothes
|
||||
# 6 Dress
|
||||
# 7 Coat
|
||||
# 8 Socks
|
||||
# 9 Pants
|
||||
# 10 Jumpsuits
|
||||
# 11 Scarf
|
||||
# 12 Skirt
|
||||
# 13 Face
|
||||
# 14 LeftArm
|
||||
# 15 RightArm
|
||||
# 16 LeftLeg
|
||||
# 17 RightLeg
|
||||
# 18 LeftShoe
|
||||
# 19 RightShoe
|
||||
head_output, tail_output = np.split(out, indices_or_sections=[1], axis=0)
|
||||
head_output = head_output.squeeze(0)
|
||||
tail_output = tail_output.squeeze(0)
|
||||
|
||||
head_output = np.stack([cv.resize(img, dsize=input_shape) for img in head_output[:, ...]])
|
||||
tail_output = np.stack([cv.resize(img, dsize=input_shape) for img in tail_output[:, ...]])
|
||||
|
||||
tail_list = np.split(tail_output, indices_or_sections=list(range(1, 20)), axis=0)
|
||||
tail_list = [arr.squeeze(0) for arr in tail_list]
|
||||
tail_list_rev = [tail_list[i] for i in range(14)]
|
||||
tail_list_rev.extend([tail_list[15], tail_list[14], tail_list[17], tail_list[16], tail_list[19], tail_list[18]])
|
||||
tail_output_rev = np.stack(tail_list_rev, axis=0)
|
||||
tail_output_rev = np.flip(tail_output_rev, axis=2)
|
||||
raw_output_all = np.mean(np.stack([head_output, tail_output_rev], axis=0), axis=0, keepdims=True)
|
||||
raw_output_all = np.argmax(raw_output_all, axis=1)
|
||||
raw_output_all = raw_output_all.transpose(1, 2, 0)
|
||||
return raw_output_all
|
||||
|
||||
|
||||
def decode_labels(gray_image):
|
||||
"""
|
||||
Colorize image according to labels
|
||||
:param gray_image: grayscale human segmentation result
|
||||
"""
|
||||
height, width, _ = gray_image.shape
|
||||
colors = [(0, 0, 0), (128, 0, 0), (255, 0, 0), (0, 85, 0), (170, 0, 51), (255, 85, 0),
|
||||
(0, 0, 85), (0, 119, 221), (85, 85, 0), (0, 85, 85), (85, 51, 0), (52, 86, 128),
|
||||
(0, 128, 0), (0, 0, 255), (51, 170, 221), (0, 255, 255),(85, 255, 170),
|
||||
(170, 255, 85), (255, 255, 0), (255, 170, 0)]
|
||||
|
||||
segm = np.stack([colors[idx] for idx in gray_image.flatten()])
|
||||
segm = segm.reshape(height, width, 3).astype(np.uint8)
|
||||
segm = cv.cvtColor(segm, cv.COLOR_BGR2RGB)
|
||||
return segm
|
||||
|
||||
|
||||
def parse_human(image, model_path, backend=cv.dnn.DNN_BACKEND_OPENCV, target=cv.dnn.DNN_TARGET_CPU):
|
||||
"""
|
||||
Prepare input for execution, run net and postprocess output to parse human.
|
||||
:param image: input image
|
||||
:param model_path: path to JPPNet model
|
||||
:param backend: name of computation backend
|
||||
:param target: name of computation target
|
||||
"""
|
||||
input = preprocess(image)
|
||||
input_h, input_w = input.shape[2:]
|
||||
output = run_net(input, model_path, backend, target)
|
||||
grayscale_out = postprocess(output, (input_w, input_h))
|
||||
segmentation = decode_labels(grayscale_out)
|
||||
return segmentation
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='Use this script to run human parsing using JPPNet',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input', '-i', required=True, help='Path to input image.')
|
||||
parser.add_argument('--model', '-m', default='lip_jppnet_384.pb', help='Path to pb model.')
|
||||
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
|
||||
help="Choose one of computation backends: "
|
||||
"%d: automatically (by default), "
|
||||
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"%d: OpenCV implementation, "
|
||||
"%d: VKCOM, "
|
||||
"%d: CUDA"% backends)
|
||||
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
|
||||
help='Choose one of target computation devices: '
|
||||
'%d: CPU target (by default), '
|
||||
'%d: OpenCL, '
|
||||
'%d: OpenCL fp16 (half-float precision), '
|
||||
'%d: NCS2 VPU, '
|
||||
'%d: HDDL VPU, '
|
||||
'%d: Vulkan, '
|
||||
'%d: CUDA, '
|
||||
'%d: CUDA fp16 (half-float preprocess)' % targets)
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
if not os.path.isfile(args.model):
|
||||
raise OSError("Model not exist")
|
||||
|
||||
image = cv.imread(args.input)
|
||||
output = parse_human(image, args.model, args.backend, args.target)
|
||||
winName = 'Deep learning human parsing in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_AUTOSIZE)
|
||||
cv.imshow(winName, output)
|
||||
cv.waitKey()
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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()
|
||||
@@ -0,0 +1,229 @@
|
||||
<!DOCTYPE html>
|
||||
|
||||
<html>
|
||||
|
||||
<head>
|
||||
<script async src="../../opencv.js" type="text/javascript"></script>
|
||||
<script src="../../utils.js" type="text/javascript"></script>
|
||||
|
||||
<script type='text/javascript'>
|
||||
var netDet = undefined, netRecogn = undefined;
|
||||
var persons = {};
|
||||
|
||||
//! [Run face detection model]
|
||||
function detectFaces(img) {
|
||||
netDet.setInputSize(new cv.Size(img.cols, img.rows));
|
||||
var out = new cv.Mat();
|
||||
netDet.detect(img, out);
|
||||
var faces = [];
|
||||
for (var i = 0, n = out.data32F.length; i < n; i += 15) {
|
||||
var left = out.data32F[i];
|
||||
var top = out.data32F[i + 1];
|
||||
var right = (out.data32F[i] + out.data32F[i + 2]);
|
||||
var bottom = (out.data32F[i + 1] + out.data32F[i + 3]);
|
||||
left = Math.min(Math.max(0, left), img.cols - 1);
|
||||
top = Math.min(Math.max(0, top), img.rows - 1);
|
||||
right = Math.min(Math.max(0, right), img.cols - 1);
|
||||
bottom = Math.min(Math.max(0, bottom), img.rows - 1);
|
||||
|
||||
if (left < right && top < bottom) {
|
||||
faces.push({
|
||||
x: left,
|
||||
y: top,
|
||||
width: right - left,
|
||||
height: bottom - top,
|
||||
x1: out.data32F[i + 4] < 0 || out.data32F[i + 4] > img.cols - 1 ? -1 : out.data32F[i + 4],
|
||||
y1: out.data32F[i + 5] < 0 || out.data32F[i + 5] > img.rows - 1 ? -1 : out.data32F[i + 5],
|
||||
x2: out.data32F[i + 6] < 0 || out.data32F[i + 6] > img.cols - 1 ? -1 : out.data32F[i + 6],
|
||||
y2: out.data32F[i + 7] < 0 || out.data32F[i + 7] > img.rows - 1 ? -1 : out.data32F[i + 7],
|
||||
x3: out.data32F[i + 8] < 0 || out.data32F[i + 8] > img.cols - 1 ? -1 : out.data32F[i + 8],
|
||||
y3: out.data32F[i + 9] < 0 || out.data32F[i + 9] > img.rows - 1 ? -1 : out.data32F[i + 9],
|
||||
x4: out.data32F[i + 10] < 0 || out.data32F[i + 10] > img.cols - 1 ? -1 : out.data32F[i + 10],
|
||||
y4: out.data32F[i + 11] < 0 || out.data32F[i + 11] > img.rows - 1 ? -1 : out.data32F[i + 11],
|
||||
x5: out.data32F[i + 12] < 0 || out.data32F[i + 12] > img.cols - 1 ? -1 : out.data32F[i + 12],
|
||||
y5: out.data32F[i + 13] < 0 || out.data32F[i + 13] > img.rows - 1 ? -1 : out.data32F[i + 13],
|
||||
confidence: out.data32F[i + 14]
|
||||
})
|
||||
}
|
||||
}
|
||||
out.delete();
|
||||
return faces;
|
||||
};
|
||||
//! [Run face detection model]
|
||||
|
||||
//! [Get 128 floating points feature vector]
|
||||
function face2vec(face) {
|
||||
var blob = cv.blobFromImage(face, 1.0, {width: 112, height: 112}, [0, 0, 0, 0], true, false)
|
||||
netRecogn.setInput(blob);
|
||||
var vec = netRecogn.forward();
|
||||
blob.delete();
|
||||
return vec;
|
||||
};
|
||||
//! [Get 128 floating points feature vector]
|
||||
|
||||
//! [Recognize]
|
||||
function recognize(face) {
|
||||
var vec = face2vec(face);
|
||||
|
||||
var bestMatchName = 'unknown';
|
||||
var bestMatchScore = 30; // Threshold for face recognition.
|
||||
for (name in persons) {
|
||||
var personVec = persons[name];
|
||||
var score = vec.dot(personVec);
|
||||
if (score > bestMatchScore) {
|
||||
bestMatchScore = score;
|
||||
bestMatchName = name;
|
||||
}
|
||||
}
|
||||
vec.delete();
|
||||
return bestMatchName;
|
||||
};
|
||||
//! [Recognize]
|
||||
|
||||
function loadModels(callback) {
|
||||
var utils = new Utils('');
|
||||
var detectModel = 'https://media.githubusercontent.com/media/opencv/opencv_zoo/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx';
|
||||
var recognModel = 'https://media.githubusercontent.com/media/opencv/opencv_zoo/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx';
|
||||
document.getElementById('status').innerHTML = 'Downloading YuNet model';
|
||||
utils.createFileFromUrl('face_detection_yunet_2023mar.onnx', detectModel, () => {
|
||||
document.getElementById('status').innerHTML = 'Downloading OpenFace model';
|
||||
utils.createFileFromUrl('face_recognition_sface_2021dec.onnx', recognModel, () => {
|
||||
document.getElementById('status').innerHTML = '';
|
||||
netDet = new cv.FaceDetectorYN("face_detection_yunet_2023mar.onnx", "", new cv.Size(320, 320), 0.9, 0.3, 5000);
|
||||
netRecogn = cv.readNet('face_recognition_sface_2021dec.onnx');
|
||||
callback();
|
||||
});
|
||||
});
|
||||
};
|
||||
|
||||
function main() {
|
||||
if(!cv.FaceDetectorYN){
|
||||
alert(`Error: This sample require OpenCV.js built with FaceDetectorYN. Please rebuild it with FaceDetectorYN or use the latest version of OpenCV.js.`);
|
||||
return;
|
||||
}
|
||||
// Create a camera object.
|
||||
var output = document.getElementById('output');
|
||||
var camera = document.createElement("video");
|
||||
camera.setAttribute("width", output.width);
|
||||
camera.setAttribute("height", output.height);
|
||||
|
||||
// Get a permission from user to use a camera.
|
||||
navigator.mediaDevices.getUserMedia({video: true, audio: false})
|
||||
.then(function(stream) {
|
||||
camera.srcObject = stream;
|
||||
camera.onloadedmetadata = function(e) {
|
||||
camera.play();
|
||||
};
|
||||
});
|
||||
|
||||
//! [Open a camera stream]
|
||||
var cap = new cv.VideoCapture(camera);
|
||||
var frame = new cv.Mat(camera.height, camera.width, cv.CV_8UC4);
|
||||
var frameBGR = new cv.Mat(camera.height, camera.width, cv.CV_8UC3);
|
||||
//! [Open a camera stream]
|
||||
|
||||
//! [Add a person]
|
||||
document.getElementById('addPersonButton').onclick = function() {
|
||||
var rects = detectFaces(frameBGR);
|
||||
if (rects.length > 0) {
|
||||
var face = frameBGR.roi(rects[0]);
|
||||
|
||||
var name = prompt('Say your name:');
|
||||
var cell = document.getElementById("targetNames").insertCell(0);
|
||||
cell.innerHTML = name;
|
||||
|
||||
persons[name] = face2vec(face).clone();
|
||||
|
||||
var canvas = document.createElement("canvas");
|
||||
canvas.setAttribute("width", 112);
|
||||
canvas.setAttribute("height", 112);
|
||||
var cell = document.getElementById("targetImgs").insertCell(0);
|
||||
cell.appendChild(canvas);
|
||||
|
||||
var faceResized = new cv.Mat(canvas.height, canvas.width, cv.CV_8UC3);
|
||||
cv.resize(face, faceResized, {width: canvas.width, height: canvas.height});
|
||||
cv.cvtColor(faceResized, faceResized, cv.COLOR_BGR2RGB);
|
||||
cv.imshow(canvas, faceResized);
|
||||
faceResized.delete();
|
||||
}
|
||||
};
|
||||
//! [Add a person]
|
||||
|
||||
//! [Define frames processing]
|
||||
var isRunning = false;
|
||||
const FPS = 30; // Target number of frames processed per second.
|
||||
function captureFrame() {
|
||||
var begin = Date.now();
|
||||
cap.read(frame); // Read a frame from camera
|
||||
cv.cvtColor(frame, frameBGR, cv.COLOR_RGBA2BGR);
|
||||
|
||||
var faces = detectFaces(frameBGR);
|
||||
faces.forEach(function(rect) {
|
||||
cv.rectangle(frame, {x: rect.x, y: rect.y}, {x: rect.x + rect.width, y: rect.y + rect.height}, [0, 255, 0, 255]);
|
||||
if(rect.x1>0 && rect.y1>0)
|
||||
cv.circle(frame, {x: rect.x1, y: rect.y1}, 2, [255, 0, 0, 255], 2)
|
||||
if(rect.x2>0 && rect.y2>0)
|
||||
cv.circle(frame, {x: rect.x2, y: rect.y2}, 2, [0, 0, 255, 255], 2)
|
||||
if(rect.x3>0 && rect.y3>0)
|
||||
cv.circle(frame, {x: rect.x3, y: rect.y3}, 2, [0, 255, 0, 255], 2)
|
||||
if(rect.x4>0 && rect.y4>0)
|
||||
cv.circle(frame, {x: rect.x4, y: rect.y4}, 2, [255, 0, 255, 255], 2)
|
||||
if(rect.x5>0 && rect.y5>0)
|
||||
cv.circle(frame, {x: rect.x5, y: rect.y5}, 2, [0, 255, 255, 255], 2)
|
||||
|
||||
var face = frameBGR.roi(rect);
|
||||
var name = recognize(face);
|
||||
cv.putText(frame, name, {x: rect.x, y: rect.y}, cv.FONT_HERSHEY_SIMPLEX, 1.0, [0, 255, 0, 255]);
|
||||
});
|
||||
|
||||
cv.imshow(output, frame);
|
||||
|
||||
// Loop this function.
|
||||
if (isRunning) {
|
||||
var delay = 1000 / FPS - (Date.now() - begin);
|
||||
setTimeout(captureFrame, delay);
|
||||
}
|
||||
};
|
||||
//! [Define frames processing]
|
||||
|
||||
document.getElementById('startStopButton').onclick = function toggle() {
|
||||
if (isRunning) {
|
||||
isRunning = false;
|
||||
document.getElementById('startStopButton').innerHTML = 'Start';
|
||||
document.getElementById('addPersonButton').disabled = true;
|
||||
} else {
|
||||
function run() {
|
||||
isRunning = true;
|
||||
captureFrame();
|
||||
document.getElementById('startStopButton').innerHTML = 'Stop';
|
||||
document.getElementById('startStopButton').disabled = false;
|
||||
document.getElementById('addPersonButton').disabled = false;
|
||||
}
|
||||
if (netDet == undefined || netRecogn == undefined) {
|
||||
document.getElementById('startStopButton').disabled = true;
|
||||
loadModels(run); // Load models and run a pipeline;
|
||||
} else {
|
||||
run();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
document.getElementById('startStopButton').disabled = false;
|
||||
};
|
||||
</script>
|
||||
|
||||
</head>
|
||||
|
||||
<body onload="cv['onRuntimeInitialized']=()=>{ main() }">
|
||||
<button id="startStopButton" type="button" disabled="true">Start</button>
|
||||
<div id="status"></div>
|
||||
<canvas id="output" width=640 height=480 style="max-width: 100%"></canvas>
|
||||
|
||||
<table>
|
||||
<tr id="targetImgs"></tr>
|
||||
<tr id="targetNames"></tr>
|
||||
</table>
|
||||
<button id="addPersonButton" type="button" disabled="true">Add a person</button>
|
||||
</body>
|
||||
|
||||
</html>
|
||||
@@ -0,0 +1,593 @@
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
import argparse
|
||||
from tqdm import tqdm
|
||||
from functools import partial
|
||||
from copy import deepcopy
|
||||
import os
|
||||
from common import *
|
||||
|
||||
## General information on how to use the sample
|
||||
|
||||
'''
|
||||
This sample proposes experimental inpainting sample using Latent Diffusion Model (LDM) for inpainting.
|
||||
Most of the script is based on the code from the official repository of the LDM model: https://github.com/CompVis/latent-diffusion
|
||||
|
||||
Current limitations of the script:
|
||||
- Slow diffusion sampling
|
||||
- Not exact reproduction of the results from the original repository (due to issues related deviation in convolution operation.
|
||||
See issue for more details: https://github.com/opencv/opencv/pull/25973)
|
||||
|
||||
LDM inpainting model was converted to ONNX graph using following steps:
|
||||
|
||||
Generate the onnx model using this [repo](https://github.com/Abdurrahheem/latent-diffusion/tree/ash/export2onnx) and follow instructions below
|
||||
|
||||
- git clone https://github.com/Abdurrahheem/latent-diffusion.git
|
||||
- cd latent-diffusion
|
||||
- conda env create -f environment.yaml
|
||||
- conda activate ldm
|
||||
- wget -O models/ldm/inpainting_big/last.ckpt https://heibox.uni-heidelberg.de/f/4d9ac7ea40c64582b7c9/?dl=1
|
||||
- python -m scripts.inpaint.py --indir data/inpainting_examples/ --outdir outputs/inpainting_results --export=True
|
||||
|
||||
2. Build opencv
|
||||
3. Run the script
|
||||
|
||||
- cd opencv/samples/dnn
|
||||
- Download models using `python download_models.py ldm_inpainting`
|
||||
- python ldm_inpainting.py
|
||||
- For more options, use python ldm_inpainting.py -h
|
||||
|
||||
After running the code you will be promted with image. You can click on left mouse button and start selecting a region you would like to be inpainted (deleted).
|
||||
Once you finish marking the region, click on left mouse button again and press esc button on your keyboard. The inpainting proccess will start.
|
||||
|
||||
Note: If you are running it on CPU it might take a large chank of time.
|
||||
Also make sure to have abount 15GB of RAM to make proccess faster (other wise swapping will kick in and everything will be slower)
|
||||
'''
|
||||
|
||||
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('--samples', '-s', type=int, help='Number of times to sample the model.', default=50)
|
||||
parser.add_argument('--mask', '-m', type=str, help='Path to mask image. If not provided, interactive mask creation will be used.', default=None)
|
||||
|
||||
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, 'ldm_inpainting', prefix="", alias="ldm_inpainting")
|
||||
add_preproc_args(args.zoo, parser, 'ldm_inpainting', prefix="encoder_", alias="ldm_inpainting")
|
||||
add_preproc_args(args.zoo, parser, 'ldm_inpainting', prefix="decoder_", alias="ldm_inpainting")
|
||||
add_preproc_args(args.zoo, parser, 'ldm_inpainting', prefix="diffusor_", alias="ldm_inpainting")
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Diffusion based image inpainting using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args()
|
||||
|
||||
stdSize = 0.7
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
imgWidth = None
|
||||
fontSize = 1.5
|
||||
fontThickness = 1
|
||||
|
||||
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 help():
|
||||
print(
|
||||
'''
|
||||
Use this script for image inpainting using OpenCV.
|
||||
|
||||
Firstly, download required models i.e. ldm_inpainting using `download_models.py ldm_inpainting` (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 ldm_inpainting.py
|
||||
'''
|
||||
)
|
||||
|
||||
def make_batch_blob(image, mask):
|
||||
|
||||
blob_image = cv.dnn.blobFromImage(image, scalefactor=args.scale, size=(args.width, args.height), mean=args.mean, swapRB=args.rgb, crop=False)
|
||||
|
||||
blob_mask = cv.dnn.blobFromImage(mask, scalefactor=args.scale, size=(args.width, args.height), mean=args.mean, swapRB=False, crop=False)
|
||||
|
||||
blob_mask = (blob_mask >= 0.5).astype(np.float32)
|
||||
masked_image = (1 - blob_mask) * blob_image
|
||||
|
||||
batch = {
|
||||
"image": blob_image,
|
||||
"mask": blob_mask,
|
||||
"masked_image": masked_image
|
||||
}
|
||||
|
||||
for k in batch:
|
||||
batch[k] = batch[k]*2.0 - 1.0
|
||||
|
||||
return batch
|
||||
|
||||
def noise_like(shape, repeat=False):
|
||||
repeat_noise = lambda: np.random.randn((1, *shape[1:])).repeat(shape[0], *((1,) * (len(shape) - 1)))
|
||||
noise = lambda: np.random.randn(*shape)
|
||||
return repeat_noise() if repeat else noise()
|
||||
|
||||
def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
|
||||
if ddim_discr_method == 'uniform':
|
||||
c = num_ddpm_timesteps // num_ddim_timesteps
|
||||
ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
|
||||
elif ddim_discr_method == 'quad':
|
||||
ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
|
||||
else:
|
||||
raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
|
||||
|
||||
# assert ddim_timesteps.shape[0] == num_ddim_timesteps
|
||||
# add one to get the final alpha values right (the ones from first scale to data during sampling)
|
||||
steps_out = ddim_timesteps + 1
|
||||
if verbose:
|
||||
print(f'Selected timesteps for ddim sampler: {steps_out}')
|
||||
return steps_out
|
||||
|
||||
def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
|
||||
# select alphas for computing the variance schedule
|
||||
alphas = alphacums[ddim_timesteps]
|
||||
alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
|
||||
|
||||
# according the the formula provided in https://arxiv.org/abs/2010.02502
|
||||
sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
|
||||
if verbose:
|
||||
print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
|
||||
print(f'For the chosen value of eta, which is {eta}, '
|
||||
f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
|
||||
return sigmas, alphas, alphas_prev
|
||||
|
||||
def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
||||
if schedule == "linear":
|
||||
betas = (
|
||||
np.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep).astype(np.float64) ** 2
|
||||
)
|
||||
|
||||
elif schedule == "cosine":
|
||||
timesteps = (
|
||||
np.arange(n_timestep + 1).astype(np.float64) / n_timestep + cosine_s
|
||||
)
|
||||
alphas = timesteps / (1 + cosine_s) * np.pi / 2
|
||||
alphas = np.cos(alphas).pow(2)
|
||||
alphas = alphas / alphas[0]
|
||||
betas = 1 - alphas[1:] / alphas[:-1]
|
||||
betas = np.clip(betas, a_min=0, a_max=0.999)
|
||||
|
||||
elif schedule == "sqrt_linear":
|
||||
betas = np.linspace(linear_start, linear_end, n_timestep).astype(np.float64)
|
||||
elif schedule == "sqrt":
|
||||
betas = np.linspace(linear_start, linear_end, n_timestep).astype(np.float64) ** 0.5
|
||||
else:
|
||||
raise ValueError(f"schedule '{schedule}' unknown.")
|
||||
return betas
|
||||
|
||||
class DDIMSampler(object):
|
||||
def __init__(self, model, schedule="linear", ddpm_num_timesteps=1000):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.ddpm_num_timesteps = ddpm_num_timesteps
|
||||
self.schedule = schedule
|
||||
|
||||
def register_buffer(self, name, attr):
|
||||
setattr(self, name, attr)
|
||||
|
||||
def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
|
||||
self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
|
||||
num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
|
||||
alphas_cumprod = self.model.alphas_cumprod
|
||||
assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
|
||||
to_numpy = partial(np.array, copy=True, dtype=np.float32)
|
||||
|
||||
self.register_buffer('betas', to_numpy(self.model.betas))
|
||||
self.register_buffer('alphas_cumprod', to_numpy(alphas_cumprod))
|
||||
self.register_buffer('alphas_cumprod_prev', to_numpy(self.model.alphas_cumprod_prev))
|
||||
|
||||
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||
self.register_buffer('sqrt_alphas_cumprod', to_numpy(np.sqrt(alphas_cumprod)))
|
||||
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_numpy(np.sqrt(1. - alphas_cumprod)))
|
||||
self.register_buffer('log_one_minus_alphas_cumprod', to_numpy(np.log(1. - alphas_cumprod)))
|
||||
self.register_buffer('sqrt_recip_alphas_cumprod', to_numpy(np.sqrt(1. / alphas_cumprod)))
|
||||
self.register_buffer('sqrt_recipm1_alphas_cumprod', to_numpy(np.sqrt(1. / alphas_cumprod - 1)))
|
||||
|
||||
# ddim sampling parameters
|
||||
ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod,
|
||||
ddim_timesteps=self.ddim_timesteps,
|
||||
eta=ddim_eta,verbose=verbose)
|
||||
self.register_buffer('ddim_sigmas', ddim_sigmas)
|
||||
self.register_buffer('ddim_alphas', ddim_alphas)
|
||||
self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
|
||||
self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
|
||||
sigmas_for_original_sampling_steps = ddim_eta * np.sqrt(
|
||||
(1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
|
||||
1 - self.alphas_cumprod / self.alphas_cumprod_prev))
|
||||
self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
|
||||
|
||||
def sample(self,
|
||||
S,
|
||||
batch_size,
|
||||
shape,
|
||||
conditioning=None,
|
||||
eta=0.,
|
||||
temperature=1.,
|
||||
verbose=True,
|
||||
x_T=None,
|
||||
log_every_t=100,
|
||||
unconditional_guidance_scale=1.,
|
||||
unconditional_conditioning=None,
|
||||
# this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
|
||||
**kwargs
|
||||
):
|
||||
if conditioning is not None:
|
||||
if isinstance(conditioning, dict):
|
||||
cbs = conditioning[list(conditioning.keys())[0]].shape[0]
|
||||
if cbs != batch_size:
|
||||
print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
|
||||
else:
|
||||
if conditioning.shape[0] != batch_size:
|
||||
print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
|
||||
|
||||
self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
|
||||
# sampling
|
||||
C, H, W = shape
|
||||
size = (batch_size, C, H, W)
|
||||
print(f'Data shape for DDIM sampling is {size}, eta {eta}')
|
||||
|
||||
samples, intermediates = self.ddim_sampling(conditioning, size,
|
||||
ddim_use_original_steps=False,
|
||||
temperature=temperature,
|
||||
x_T=x_T,
|
||||
log_every_t=log_every_t,
|
||||
unconditional_guidance_scale=unconditional_guidance_scale,
|
||||
unconditional_conditioning=unconditional_conditioning,
|
||||
)
|
||||
return samples, intermediates
|
||||
|
||||
def ddim_sampling(self, cond, shape,
|
||||
x_T=None, ddim_use_original_steps=False,
|
||||
timesteps=None,log_every_t=100, temperature=1.,
|
||||
unconditional_guidance_scale=1., unconditional_conditioning=None,):
|
||||
b = shape[0]
|
||||
if x_T is None:
|
||||
img = np.random.randn(*shape)
|
||||
else:
|
||||
img = x_T
|
||||
|
||||
if timesteps is None:
|
||||
timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
|
||||
elif timesteps is not None and not ddim_use_original_steps:
|
||||
subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
|
||||
timesteps = self.ddim_timesteps[:subset_end]
|
||||
|
||||
intermediates = {'x_inter': [img], 'pred_x0': [img]}
|
||||
time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
|
||||
total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
|
||||
print(f"Running DDIM Sampling with {total_steps} timesteps")
|
||||
|
||||
iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
|
||||
|
||||
for i, step in enumerate(iterator):
|
||||
index = total_steps - i - 1
|
||||
ts = np.full((b, ), step, dtype=np.int64)
|
||||
|
||||
outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
|
||||
temperature=temperature, unconditional_guidance_scale=unconditional_guidance_scale,
|
||||
unconditional_conditioning=unconditional_conditioning)
|
||||
img, pred_x0 = outs
|
||||
if index % log_every_t == 0 or index == total_steps - 1:
|
||||
intermediates['x_inter'].append(img)
|
||||
intermediates['pred_x0'].append(pred_x0)
|
||||
|
||||
return img, intermediates
|
||||
|
||||
def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False,
|
||||
temperature=1., unconditional_guidance_scale=1., unconditional_conditioning=None):
|
||||
b = x.shape[0]
|
||||
if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
|
||||
e_t = self.model.apply_model(x, t, c)
|
||||
|
||||
alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
|
||||
alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
|
||||
sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
|
||||
sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
|
||||
# select parameters corresponding to the currently considered timestep
|
||||
a_t = np.full((b, 1, 1, 1), alphas[index])
|
||||
a_prev = np.full((b, 1, 1, 1), alphas_prev[index])
|
||||
sigma_t = np.full((b, 1, 1, 1), sigmas[index])
|
||||
sqrt_one_minus_at = np.full((b, 1, 1, 1), sqrt_one_minus_alphas[index])
|
||||
|
||||
# current prediction for x_0
|
||||
pred_x0 = (x - sqrt_one_minus_at * e_t) / np.sqrt(a_t)
|
||||
# direction pointing to x_t
|
||||
dir_xt = np.sqrt(1. - a_prev - sigma_t**2) * e_t
|
||||
noise = sigma_t * noise_like(x.shape, repeat_noise) * temperature
|
||||
x_prev = np.sqrt(a_prev) * pred_x0 + dir_xt + noise
|
||||
return x_prev, pred_x0
|
||||
|
||||
|
||||
class DDIMInpainter(object):
|
||||
def __init__(self,
|
||||
args,
|
||||
v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
|
||||
parameterization="eps", # all assuming fixed variance schedules
|
||||
linear_start=0.0015,
|
||||
linear_end=0.0205,
|
||||
conditioning_key="concat",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.v_posterior = v_posterior
|
||||
self.parameterization = parameterization
|
||||
self.conditioning_key = conditioning_key
|
||||
self.register_schedule(linear_start=linear_start, linear_end=linear_end)
|
||||
|
||||
# Initialize models using provided paths or download if necessary
|
||||
encoder_path = findModel(args.encoder_model, args.encoder_sha1)
|
||||
decoder_path = findModel(args.decoder_model, args.decoder_sha1)
|
||||
diffusor_path = findModel(args.diffusor_model, args.diffusor_sha1)
|
||||
|
||||
engine = cv.dnn.ENGINE_AUTO
|
||||
if args.backend != "default" or args.target != "cpu":
|
||||
engine = cv.dnn.ENGINE_CLASSIC
|
||||
|
||||
self.encoder = cv.dnn.readNet(encoder_path, "", "", engine)
|
||||
self.diffusor = cv.dnn.readNet(diffusor_path, "", "", engine)
|
||||
self.decoder = cv.dnn.readNet(decoder_path, "", "", engine)
|
||||
self.sampler = DDIMSampler(self, ddpm_num_timesteps=self.num_timesteps)
|
||||
self.set_backend(backend=get_backend_id(args.backend), target=get_target_id(args.target))
|
||||
|
||||
def set_backend(self, backend=cv.dnn.DNN_BACKEND_DEFAULT, target=cv.dnn.DNN_TARGET_CPU):
|
||||
self.encoder.setPreferableBackend(backend)
|
||||
self.encoder.setPreferableTarget(target)
|
||||
|
||||
self.decoder.setPreferableBackend(backend)
|
||||
self.decoder.setPreferableTarget(target)
|
||||
|
||||
self.diffusor.setPreferableBackend(backend)
|
||||
self.diffusor.setPreferableTarget(target)
|
||||
|
||||
def apply_diffusor(self, x, timestep, cond):
|
||||
x = np.concatenate([x, cond], axis=1)
|
||||
x = cv.Mat(x.astype(np.float32))
|
||||
timestep = cv.Mat(timestep.astype(np.int64))
|
||||
names = ["xc, t", "timesteps"]
|
||||
self.diffusor.setInputsNames(names)
|
||||
self.diffusor.setInput(x, names[0])
|
||||
self.diffusor.setInput(timestep, names[1])
|
||||
output = self.diffusor.forward()
|
||||
|
||||
return output
|
||||
|
||||
def register_buffer(self, name, attr):
|
||||
setattr(self, name, attr)
|
||||
|
||||
def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
|
||||
linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
||||
if given_betas is not None:
|
||||
betas = given_betas
|
||||
else:
|
||||
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
|
||||
cosine_s=cosine_s)
|
||||
alphas = 1. - betas
|
||||
alphas_cumprod = np.cumprod(alphas, axis=0)
|
||||
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
|
||||
|
||||
timesteps, = betas.shape
|
||||
self.num_timesteps = int(timesteps)
|
||||
self.linear_start = linear_start
|
||||
self.linear_end = linear_end
|
||||
assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
|
||||
|
||||
to_numpy = partial(np.array, dtype=np.float32)
|
||||
|
||||
self.register_buffer('betas', to_numpy(betas))
|
||||
self.register_buffer('alphas_cumprod', to_numpy(alphas_cumprod))
|
||||
self.register_buffer('alphas_cumprod_prev', to_numpy(alphas_cumprod_prev))
|
||||
|
||||
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||
self.register_buffer('sqrt_alphas_cumprod', to_numpy(np.sqrt(alphas_cumprod)))
|
||||
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_numpy(np.sqrt(1. - alphas_cumprod)))
|
||||
self.register_buffer('log_one_minus_alphas_cumprod', to_numpy(np.log(1. - alphas_cumprod)))
|
||||
self.register_buffer('sqrt_recip_alphas_cumprod', to_numpy(np.sqrt(1. / alphas_cumprod)))
|
||||
self.register_buffer('sqrt_recipm1_alphas_cumprod', to_numpy(np.sqrt(1. / alphas_cumprod - 1)))
|
||||
|
||||
# calculations for posterior q(x_{t-1} | x_t, x_0)
|
||||
posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
|
||||
1. - alphas_cumprod) + self.v_posterior * betas
|
||||
# above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
|
||||
self.register_buffer('posterior_variance', to_numpy(posterior_variance))
|
||||
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
|
||||
self.register_buffer('posterior_log_variance_clipped', to_numpy(np.log(np.maximum(posterior_variance, 1e-20))))
|
||||
self.register_buffer('posterior_mean_coef1', to_numpy(
|
||||
betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
|
||||
self.register_buffer('posterior_mean_coef2', to_numpy(
|
||||
(1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
|
||||
if self.parameterization == "eps":
|
||||
lvlb_weights = self.betas ** 2 / (
|
||||
2 * self.posterior_variance * to_numpy(alphas) * (1 - self.alphas_cumprod))
|
||||
elif self.parameterization == "x0":
|
||||
lvlb_weights = 0.5 * np.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
|
||||
else:
|
||||
raise NotImplementedError("mu not supported")
|
||||
# TODO how to choose this term
|
||||
lvlb_weights[0] = lvlb_weights[1]
|
||||
self.register_buffer('lvlb_weights', lvlb_weights)
|
||||
assert not np.isnan(self.lvlb_weights).all()
|
||||
|
||||
def apply_model(self, x_noisy, t, cond, return_ids=False):
|
||||
if isinstance(cond, dict):
|
||||
# hybrid case, cond is exptected to be a dict
|
||||
pass
|
||||
else:
|
||||
# if not isinstance(cond, list):
|
||||
# cond = [cond]
|
||||
key = 'c_concat' if self.conditioning_key == 'concat' else 'c_crossattn'
|
||||
cond = {key: cond}
|
||||
|
||||
x_recon = self.apply_diffusor(x_noisy, t, cond['c_concat'])
|
||||
if isinstance(x_recon, tuple) and not return_ids:
|
||||
return x_recon[0]
|
||||
else:
|
||||
return x_recon
|
||||
|
||||
def inpaint(self, image : np.ndarray, mask : np.ndarray, S : int = 50) -> np.ndarray:
|
||||
inpainted = self(image, mask, S)
|
||||
return np.squeeze(inpainted)
|
||||
|
||||
def __call__(self, image : np.ndarray, mask : np.ndarray, S : int = 50) -> np.ndarray:
|
||||
|
||||
# Encode the image and mask
|
||||
self.encoder.setInput(image)
|
||||
c = self.encoder.forward()
|
||||
cc = cv.resize(np.squeeze(mask), dsize=(c.shape[3], c.shape[2]), interpolation=cv.INTER_NEAREST) #TODO:check for correcteness of intepolation
|
||||
cc = cc[None,None]
|
||||
c = np.concatenate([c, cc], axis=1)
|
||||
|
||||
shape = (c.shape[1] - 1,) + c.shape[2:]
|
||||
# Sample from the model
|
||||
samples_ddim, _ = self.sampler.sample(
|
||||
S=S,
|
||||
conditioning=c,
|
||||
batch_size=c.shape[0],
|
||||
shape=shape,
|
||||
verbose=False)
|
||||
|
||||
## Decode the sample
|
||||
samples_ddim = samples_ddim.astype(np.float32)
|
||||
samples_ddim = cv.Mat(samples_ddim)
|
||||
self.decoder.setInput(samples_ddim)
|
||||
x_samples_ddim = self.decoder.forward()
|
||||
|
||||
image = np.clip((image + 1.0) / 2.0, a_min=0.0, a_max=1.0)
|
||||
mask = np.clip((mask + 1.0) / 2.0, a_min=0.0, a_max=1.0)
|
||||
predicted_image = np.clip((x_samples_ddim + 1.0) / 2.0, a_min=0.0, a_max=1.0)
|
||||
|
||||
inpainted = (1 - mask) * image + mask * predicted_image
|
||||
inpainted = np.transpose(inpainted, (0, 2, 3, 1)) * 255
|
||||
|
||||
return inpainted
|
||||
|
||||
def create_mask(img):
|
||||
drawing = False # True if the mouse is pressed
|
||||
brush_size = 20
|
||||
|
||||
# Mouse callback function
|
||||
def draw_circle(event, x, y, flags, param):
|
||||
nonlocal drawing, brush_size
|
||||
|
||||
if event == cv.EVENT_LBUTTONDOWN:
|
||||
drawing = True
|
||||
elif event == cv.EVENT_MOUSEMOVE:
|
||||
if drawing:
|
||||
cv.circle(mask, (x, y), brush_size, (255), thickness=-1)
|
||||
elif event == cv.EVENT_LBUTTONUP:
|
||||
drawing = False
|
||||
|
||||
|
||||
# Create window with instructions
|
||||
window_name = 'Draw Mask'
|
||||
cv.namedWindow(window_name)
|
||||
cv.setMouseCallback(window_name, draw_circle)
|
||||
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
|
||||
temp_image = img.copy()
|
||||
overlay = img.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)
|
||||
|
||||
mask = np.zeros((img.shape[0], img.shape[1]), np.uint8)
|
||||
display_img = temp_image.copy()
|
||||
while True:
|
||||
display_img[mask > 0] = [255, 255, 255]
|
||||
cv.imshow(window_name, display_img)
|
||||
# Create a copy of the image to show instructions
|
||||
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 = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
|
||||
display_img = temp_image.copy()
|
||||
print(f"Mask cleared")
|
||||
elif key == ord(' '): # Press space bar to finish drawing
|
||||
break
|
||||
elif key == 27:
|
||||
exit()
|
||||
|
||||
cv.destroyAllWindows()
|
||||
return mask
|
||||
|
||||
def prepare_input(args, image):
|
||||
if args.mask:
|
||||
mask = cv.imread(args.mask, cv.IMREAD_GRAYSCALE)
|
||||
if mask is None:
|
||||
raise ValueError(f"Could not read mask file: {args.mask}")
|
||||
if mask.shape[:2] != image.shape[:2]:
|
||||
mask = cv.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv.INTER_NEAREST)
|
||||
else:
|
||||
mask = create_mask(deepcopy(image))
|
||||
|
||||
batch = make_batch_blob(image, mask)
|
||||
return batch
|
||||
|
||||
def main(args):
|
||||
global imgWidth, fontSize, fontThickness
|
||||
keyboard_shorcuts()
|
||||
|
||||
image = cv.imread(findFile(args.input))
|
||||
imgWidth = min(image.shape[:2])
|
||||
fontSize = min(1.5, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(1,(stdWeight*imgWidth)//stdImgSize)
|
||||
aspect_ratio = image.shape[0]/image.shape[1]
|
||||
height = int(args.width*aspect_ratio)
|
||||
|
||||
batch = prepare_input(args, image)
|
||||
|
||||
model = DDIMInpainter(args)
|
||||
result = model.inpaint(batch["masked_image"], batch["mask"], S=args.samples)
|
||||
|
||||
result = result.astype(np.uint8)
|
||||
result = cv.resize(result, (args.width, height))
|
||||
result = cv.cvtColor(result, cv.COLOR_RGB2BGR)
|
||||
cv.imshow("Inpainted Image", result)
|
||||
cv.waitKey(0)
|
||||
cv.destroyAllWindows()
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args_parser()
|
||||
main(args)
|
||||
@@ -0,0 +1,158 @@
|
||||
'''
|
||||
Mask R-CNN
|
||||
This is an example of using Mask R-CNN for object detection and instance segmentation.
|
||||
|
||||
NOTE regarding OpenCV 5.0+:
|
||||
The default model configuration (.pbtxt) used in this sample relies on retrieving
|
||||
intermediate layers (e.g., 'detection_out_final'). OpenCV 5.0 introduces stricter
|
||||
graph optimization which may prune intermediate layers not explicitly registered as outputs.
|
||||
If you encounter an error such as "the number of requested and actual outputs must be the same",
|
||||
please note that the provided .pbtxt may need to be updated to explicitly declare
|
||||
'detection_out_final' as an output node.
|
||||
'''
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
parser = argparse.ArgumentParser(description=
|
||||
'Use this script to run Mask-RCNN object detection and semantic '
|
||||
'segmentation network from TensorFlow Object Detection API.')
|
||||
parser.add_argument('--input', help='Path to input image or video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--model', required=True, help='Path to a .pb file with weights.')
|
||||
parser.add_argument('--config', required=True, help='Path to a .pxtxt file contains network configuration.')
|
||||
parser.add_argument('--classes', help='Optional path to a text file with names of classes.')
|
||||
parser.add_argument('--colors', help='Optional path to a text file with colors for an every class. '
|
||||
'An every color is represented with three values from 0 to 255 in BGR channels order.')
|
||||
parser.add_argument('--width', type=int, default=800,
|
||||
help='Preprocess input image by resizing to a specific width.')
|
||||
parser.add_argument('--height', type=int, default=800,
|
||||
help='Preprocess input image by resizing to a specific height.')
|
||||
parser.add_argument('--thr', type=float, default=0.5, help='Confidence threshold')
|
||||
args = parser.parse_args()
|
||||
|
||||
np.random.seed(324)
|
||||
|
||||
# Load names of classes
|
||||
classes = None
|
||||
if args.classes:
|
||||
with open(args.classes, 'rt') as f:
|
||||
classes = f.read().rstrip('\n').split('\n')
|
||||
|
||||
# Load colors
|
||||
colors = None
|
||||
if args.colors:
|
||||
with open(args.colors, 'rt') as f:
|
||||
colors = [np.array(color.split(' '), np.uint8) for color in f.read().rstrip('\n').split('\n')]
|
||||
|
||||
legend = None
|
||||
def showLegend(classes):
|
||||
global legend
|
||||
if not classes is None and legend is None:
|
||||
blockHeight = 30
|
||||
assert(len(classes) == len(colors))
|
||||
|
||||
legend = np.zeros((blockHeight * len(colors), 200, 3), np.uint8)
|
||||
for i in range(len(classes)):
|
||||
block = legend[i * blockHeight:(i + 1) * blockHeight]
|
||||
block[:,:] = colors[i]
|
||||
cv.putText(block, classes[i], (0, blockHeight//2), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255))
|
||||
|
||||
cv.namedWindow('Legend', cv.WINDOW_NORMAL)
|
||||
cv.imshow('Legend', legend)
|
||||
classes = None
|
||||
|
||||
|
||||
def drawBox(frame, classId, conf, left, top, right, bottom):
|
||||
# Draw a bounding box.
|
||||
cv.rectangle(frame, (left, top), (right, bottom), (0, 255, 0))
|
||||
|
||||
label = '%.2f' % conf
|
||||
|
||||
# Print a label of class.
|
||||
if classes:
|
||||
assert(classId < len(classes))
|
||||
label = '%s: %s' % (classes[classId], label)
|
||||
|
||||
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||
top = max(top, labelSize[1])
|
||||
cv.rectangle(frame, (left, top - labelSize[1]), (left + labelSize[0], top + baseLine), (255, 255, 255), cv.FILLED)
|
||||
cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
|
||||
|
||||
|
||||
# Load a network
|
||||
net = cv.dnn.readNet(cv.samples.findFile(args.model), cv.samples.findFile(args.config))
|
||||
net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
|
||||
|
||||
winName = 'Mask-RCNN in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_NORMAL)
|
||||
|
||||
cap = cv.VideoCapture(cv.samples.findFileOrKeep(args.input) if args.input else 0)
|
||||
legend = None
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
cv.waitKey()
|
||||
break
|
||||
|
||||
frameH = frame.shape[0]
|
||||
frameW = frame.shape[1]
|
||||
|
||||
# Create a 4D blob from a frame.
|
||||
blob = cv.dnn.blobFromImage(frame, size=(args.width, args.height), swapRB=True, crop=False)
|
||||
|
||||
# Run a model
|
||||
net.setInput(blob)
|
||||
|
||||
# NOTE: In OpenCV 5.0, requesting 'detection_out_final' will fail if the .pbtxt
|
||||
# does not register it as an output. See file header for details.
|
||||
t0 = cv.getTickCount()
|
||||
boxes, masks = net.forward(['detection_out_final', 'detection_masks'])
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
|
||||
numClasses = masks.shape[1]
|
||||
numDetections = boxes.shape[2]
|
||||
|
||||
# Draw segmentation
|
||||
if not colors:
|
||||
# Generate colors
|
||||
colors = [np.array([0, 0, 0], np.uint8)]
|
||||
for i in range(1, numClasses + 1):
|
||||
colors.append((colors[i - 1] + np.random.randint(0, 256, [3], np.uint8)) / 2)
|
||||
del colors[0]
|
||||
|
||||
boxesToDraw = []
|
||||
for i in range(numDetections):
|
||||
box = boxes[0, 0, i]
|
||||
mask = masks[i]
|
||||
score = box[2]
|
||||
if score > args.thr:
|
||||
classId = int(box[1])
|
||||
left = int(frameW * box[3])
|
||||
top = int(frameH * box[4])
|
||||
right = int(frameW * box[5])
|
||||
bottom = int(frameH * box[6])
|
||||
|
||||
left = max(0, min(left, frameW - 1))
|
||||
top = max(0, min(top, frameH - 1))
|
||||
right = max(0, min(right, frameW - 1))
|
||||
bottom = max(0, min(bottom, frameH - 1))
|
||||
|
||||
boxesToDraw.append([frame, classId, score, left, top, right, bottom])
|
||||
|
||||
classMask = mask[classId]
|
||||
classMask = cv.resize(classMask, (right - left + 1, bottom - top + 1))
|
||||
mask = (classMask > 0.5)
|
||||
|
||||
roi = frame[top:bottom+1, left:right+1][mask]
|
||||
frame[top:bottom+1, left:right+1][mask] = (0.7 * colors[classId] + 0.3 * roi).astype(np.uint8)
|
||||
|
||||
for box in boxesToDraw:
|
||||
drawBox(*box)
|
||||
|
||||
# Put efficiency information.
|
||||
label = 'Inference time: %.2f ms' % (t * 1000.0)
|
||||
cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
|
||||
|
||||
showLegend(classes)
|
||||
|
||||
cv.imshow(winName, frame)
|
||||
@@ -0,0 +1,133 @@
|
||||
from __future__ import print_function
|
||||
# Script to evaluate MobileNet-SSD object detection model trained in TensorFlow
|
||||
# using both TensorFlow and OpenCV. Example:
|
||||
#
|
||||
# python mobilenet_ssd_accuracy.py \
|
||||
# --weights=frozen_inference_graph.pb \
|
||||
# --prototxt=ssd_mobilenet_v1_coco.pbtxt \
|
||||
# --images=val2017 \
|
||||
# --annotations=annotations/instances_val2017.json
|
||||
#
|
||||
# Tested on COCO 2017 object detection dataset, http://cocodataset.org/#download
|
||||
import os
|
||||
import cv2 as cv
|
||||
import json
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Evaluate MobileNet-SSD model using both TensorFlow and OpenCV. '
|
||||
'COCO evaluation framework is required: http://cocodataset.org')
|
||||
parser.add_argument('--weights', required=True,
|
||||
help='Path to frozen_inference_graph.pb of MobileNet-SSD model. '
|
||||
'Download it from http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_11_06_2017.tar.gz')
|
||||
parser.add_argument('--prototxt', help='Path to ssd_mobilenet_v1_coco.pbtxt from opencv_extra.', required=True)
|
||||
parser.add_argument('--images', help='Path to COCO validation images directory.', required=True)
|
||||
parser.add_argument('--annotations', help='Path to COCO annotations file.', required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
### Get OpenCV predictions #####################################################
|
||||
net = cv.dnn.readNetFromTensorflow(cv.samples.findFile(args.weights), cv.samples.findFile(args.prototxt))
|
||||
net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
|
||||
|
||||
detections = []
|
||||
for imgName in os.listdir(args.images):
|
||||
inp = cv.imread(cv.samples.findFile(os.path.join(args.images, imgName)))
|
||||
rows = inp.shape[0]
|
||||
cols = inp.shape[1]
|
||||
inp = cv.resize(inp, (300, 300))
|
||||
|
||||
net.setInput(cv.dnn.blobFromImage(inp, 1.0/127.5, (300, 300), (127.5, 127.5, 127.5), True))
|
||||
out = net.forward()
|
||||
|
||||
for i in range(out.shape[2]):
|
||||
score = float(out[0, 0, i, 2])
|
||||
# Confidence threshold is in prototxt.
|
||||
classId = int(out[0, 0, i, 1])
|
||||
|
||||
x = out[0, 0, i, 3] * cols
|
||||
y = out[0, 0, i, 4] * rows
|
||||
w = out[0, 0, i, 5] * cols - x
|
||||
h = out[0, 0, i, 6] * rows - y
|
||||
detections.append({
|
||||
"image_id": int(imgName.rstrip('0')[:imgName.rfind('.')]),
|
||||
"category_id": classId,
|
||||
"bbox": [x, y, w, h],
|
||||
"score": score
|
||||
})
|
||||
|
||||
with open('cv_result.json', 'wt') as f:
|
||||
json.dump(detections, f)
|
||||
|
||||
### Get TensorFlow predictions #################################################
|
||||
import tensorflow as tf
|
||||
|
||||
with tf.gfile.FastGFile(args.weights) as f:
|
||||
# Load the model
|
||||
graph_def = tf.GraphDef()
|
||||
graph_def.ParseFromString(f.read())
|
||||
|
||||
with tf.Session() as sess:
|
||||
# Restore session
|
||||
sess.graph.as_default()
|
||||
tf.import_graph_def(graph_def, name='')
|
||||
|
||||
detections = []
|
||||
for imgName in os.listdir(args.images):
|
||||
inp = cv.imread(os.path.join(args.images, imgName))
|
||||
rows = inp.shape[0]
|
||||
cols = inp.shape[1]
|
||||
inp = cv.resize(inp, (300, 300))
|
||||
inp = inp[:, :, [2, 1, 0]] # BGR2RGB
|
||||
out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
|
||||
sess.graph.get_tensor_by_name('detection_scores:0'),
|
||||
sess.graph.get_tensor_by_name('detection_boxes:0'),
|
||||
sess.graph.get_tensor_by_name('detection_classes:0')],
|
||||
feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
|
||||
num_detections = int(out[0][0])
|
||||
for i in range(num_detections):
|
||||
classId = int(out[3][0][i])
|
||||
score = float(out[1][0][i])
|
||||
bbox = [float(v) for v in out[2][0][i]]
|
||||
if score > 0.01:
|
||||
x = bbox[1] * cols
|
||||
y = bbox[0] * rows
|
||||
w = bbox[3] * cols - x
|
||||
h = bbox[2] * rows - y
|
||||
detections.append({
|
||||
"image_id": int(imgName.rstrip('0')[:imgName.rfind('.')]),
|
||||
"category_id": classId,
|
||||
"bbox": [x, y, w, h],
|
||||
"score": score
|
||||
})
|
||||
|
||||
with open('tf_result.json', 'wt') as f:
|
||||
json.dump(detections, f)
|
||||
|
||||
### Evaluation part ############################################################
|
||||
|
||||
# %matplotlib inline
|
||||
import matplotlib.pyplot as plt
|
||||
from pycocotools.coco import COCO
|
||||
from pycocotools.cocoeval import COCOeval
|
||||
import numpy as np
|
||||
import skimage.io as io
|
||||
import pylab
|
||||
pylab.rcParams['figure.figsize'] = (10.0, 8.0)
|
||||
|
||||
annType = ['segm','bbox','keypoints']
|
||||
annType = annType[1] #specify type here
|
||||
prefix = 'person_keypoints' if annType=='keypoints' else 'instances'
|
||||
print('Running demo for *%s* results.'%(annType))
|
||||
|
||||
#initialize COCO ground truth api
|
||||
cocoGt=COCO(args.annotations)
|
||||
|
||||
#initialize COCO detections api
|
||||
for resFile in ['tf_result.json', 'cv_result.json']:
|
||||
print(resFile)
|
||||
cocoDt=cocoGt.loadRes(resFile)
|
||||
|
||||
cocoEval = COCOeval(cocoGt,cocoDt,annType)
|
||||
cocoEval.evaluate()
|
||||
cocoEval.accumulate()
|
||||
cocoEval.summarize()
|
||||
@@ -0,0 +1,517 @@
|
||||
%YAML 1.0
|
||||
---
|
||||
################################################################################
|
||||
# Object detection models.
|
||||
################################################################################
|
||||
|
||||
# YOLOv8 object detection family from ultralytics (https://github.com/ultralytics/ultralytics)
|
||||
# Might be used for all YOLOv8n YOLOv8s YOLOv8m YOLOv8l and YOLOv8x
|
||||
yolov8x:
|
||||
load_info:
|
||||
url: "https://huggingface.co/cabelo/yolov8/resolve/main/yolov8x.onnx?download=true"
|
||||
sha1: "462f15d668c046d38e27d3df01fe8142dd004cb4"
|
||||
model: "yolov8x.onnx"
|
||||
mean: 0.0
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov8"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov8s:
|
||||
load_info:
|
||||
url: "https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov8s.onnx"
|
||||
sha1: "82cd83984396fe929909ecb58212b0e86d0904b1"
|
||||
model: "yolov8s.onnx"
|
||||
mean: 0.0
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov8"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov8:
|
||||
load_info:
|
||||
url: "https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov8n.onnx"
|
||||
sha1: "68f864475d06e2ec4037181052739f268eeac38d"
|
||||
model: "yolov8n.onnx"
|
||||
mean: 0.0
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov8"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov8m:
|
||||
load_info:
|
||||
url: "https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov8m.onnx"
|
||||
sha1: "656ffeb4f3b067bc30df956728b5f9c61a4cb090"
|
||||
model: "yolov8m.onnx"
|
||||
mean: 0.0
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov8"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov8l:
|
||||
load_info:
|
||||
url: "https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov8l.onnx"
|
||||
sha1: "462df53ca3a85d110bf6be7fc2e2bb1277124395"
|
||||
model: "yolov8l.onnx"
|
||||
mean: 0.0
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov8"
|
||||
sample: "object_detection"
|
||||
|
||||
# YOLOv5 object detection family from ultralytics (https://github.com/ultralytics/ultralytics)
|
||||
# Might be used for all YOLOv5n YOLOv5s YOLOv5m YOLOv5l and YOLOv5x
|
||||
|
||||
yolov5l:
|
||||
load_info:
|
||||
url: "https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov5l.onnx"
|
||||
sha1: "9de7e54c524b7fe7577bbd4cdbbdaed53375c8f1"
|
||||
model: "yolov5l.onnx"
|
||||
mean: 0.0
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov5"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov4:
|
||||
load_info:
|
||||
url: "https://huggingface.co/opencv/opencv_contribution/resolve/main/yolov4/yolov4.onnx"
|
||||
sha1: "68df7133bef095d79531ad62d79295d82614de3b"
|
||||
model: "yolov4.onnx"
|
||||
mean: [0, 0, 0]
|
||||
scale: 0.00392
|
||||
width: 608
|
||||
height: 608
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov4"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov4-tiny:
|
||||
load_info:
|
||||
url: "https://huggingface.co/opencv/opencv_contribution/resolve/main/yolov4/yolov4-tiny.onnx"
|
||||
sha1: "158a74e9c6da57f5e4161c5dfc1ab592f47d958a"
|
||||
model: "yolov4-tiny.onnx"
|
||||
mean: [0, 0, 0]
|
||||
scale: 0.00392
|
||||
width: 416
|
||||
height: 416
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov4"
|
||||
sample: "object_detection"
|
||||
|
||||
yolov3:
|
||||
load_info:
|
||||
url: "https://huggingface.co/qualcomm/Yolo-v3/resolve/226ada6de9dcb32eebad7f74bf526714e2af6136/Yolo-v3.onnx"
|
||||
sha1: "c37641ddf05cfe133efd4b66832f269d95f523cf"
|
||||
model: "yolov3.onnx"
|
||||
mean: [0, 0, 0]
|
||||
scale: 0.00392
|
||||
width: 640
|
||||
height: 640
|
||||
rgb: true
|
||||
labels: "object_detection_classes_yolo.txt"
|
||||
postprocessing: "yolov4"
|
||||
sample: "object_detection"
|
||||
|
||||
# Caffe implementation of SSD model from https://github.com/PINTO0309/MobileNet-SSD-RealSense
|
||||
ssd_caffe:
|
||||
load_info:
|
||||
url: "https://github.com/PINTO0309/MobileNet-SSD-RealSense/raw/refs/heads/master/caffemodel/MobileNetSSD/MobileNetSSD_deploy.caffemodel"
|
||||
sha1: "994d30a8afaa9e754d17d2373b2d62a7dfbaaf7a"
|
||||
model: "MobileNetSSD_deploy.caffemodel"
|
||||
config_load_info:
|
||||
url: "https://github.com/PINTO0309/MobileNet-SSD-RealSense/raw/refs/heads/master/caffemodel/MobileNetSSD/MobileNetSSD_deploy.prototxt"
|
||||
sha1: "25c8404cecdef638c2bd9ac7f3b46a8b96897deb"
|
||||
config: "MobileNetSSD_deploy.prototxt"
|
||||
mean: [127.5, 127.5, 127.5]
|
||||
scale: 0.007843
|
||||
width: 300
|
||||
height: 300
|
||||
rgb: false
|
||||
labels: "object_detection_classes_pascal_voc.txt"
|
||||
postprocessing: "ssd"
|
||||
sample: "object_detection"
|
||||
|
||||
# TensorFlow implementation of SSD model from https://github.com/tensorflow/models/tree/master/research/object_detection
|
||||
ssd_tf:
|
||||
load_info:
|
||||
url: "http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2017_11_17.tar.gz"
|
||||
sha1: "9e4bcdd98f4c6572747679e4ce570de4f03a70e2"
|
||||
download_sha: "6157ddb6da55db2da89dd561eceb7f944928e317"
|
||||
download_name: "ssd_mobilenet_v1_coco_2017_11_17.tar.gz"
|
||||
member: "ssd_mobilenet_v1_coco_2017_11_17/frozen_inference_graph.pb"
|
||||
model: "ssd_mobilenet_v1_coco_2017_11_17.pb"
|
||||
config_load_info:
|
||||
url: "https://github.com/opencv/opencv_extra/raw/refs/heads/4.x/testdata/dnn/ssd_mobilenet_v1_coco_2017_11_17.pbtxt"
|
||||
sha1: "c7cf985ce0a4a8953daaa4b8cacdd3c8e31437a6"
|
||||
config: "ssd_mobilenet_v1_coco_2017_11_17.pbtxt"
|
||||
mean: [0, 0, 0]
|
||||
scale: 1.0
|
||||
width: 300
|
||||
height: 300
|
||||
rgb: true
|
||||
labels: "object_detection_classes_coco.txt"
|
||||
postprocessing: "ssd"
|
||||
sample: "object_detection"
|
||||
|
||||
# TensorFlow implementation of Faster-RCNN model from https://github.com/tensorflow/models/tree/master/research/object_detection
|
||||
faster_rcnn_tf:
|
||||
load_info:
|
||||
url: "http://download.tensorflow.org/models/object_detection/faster_rcnn_inception_v2_coco_2018_01_28.tar.gz"
|
||||
sha1: "f2e4bf386b9bb3e25ddfcbbd382c20f417e444f3"
|
||||
download_sha: "c710f25e5c6a3ce85fe793d5bf266d581ab1c230"
|
||||
download_name: "faster_rcnn_inception_v2_coco_2018_01_28.tar.gz"
|
||||
member: "faster_rcnn_inception_v2_coco_2018_01_28/frozen_inference_graph.pb"
|
||||
model: "faster_rcnn_inception_v2_coco_2018_01_28.pb"
|
||||
config_load_info:
|
||||
url: "https://github.com/opencv/opencv_extra/raw/refs/heads/4.x/testdata/dnn/faster_rcnn_inception_v2_coco_2018_01_28.pbtxt"
|
||||
sha1: "059ee437fb4d6f82a6f1d2b3c7a8dd54c107687f"
|
||||
config: "faster_rcnn_inception_v2_coco_2018_01_28.pbtxt"
|
||||
mean: [0, 0, 0]
|
||||
scale: 1.0
|
||||
width: 800
|
||||
height: 600
|
||||
rgb: true
|
||||
postprocessing: "ssd"
|
||||
sample: "object_detection"
|
||||
|
||||
################################################################################
|
||||
# Image classification models.
|
||||
################################################################################
|
||||
|
||||
squeezenet:
|
||||
load_info:
|
||||
url: "https://github.com/onnx/models/raw/main/validated/vision/classification/squeezenet/model/squeezenet1.1-7.onnx?download="
|
||||
sha1: "ec31942d17715941bb9b81f3a91dc59def9236be"
|
||||
model: "squeezenet1.1-7.onnx"
|
||||
mean: [0.485, 0.456, 0.406]
|
||||
std: [0.229, 0.224, 0.225]
|
||||
scale: 0.003921
|
||||
width: 224
|
||||
height: 224
|
||||
rgb: true
|
||||
labels: "classification_classes_ILSVRC2012.txt"
|
||||
sample: "classification"
|
||||
|
||||
googlenet:
|
||||
load_info:
|
||||
url: "https://github.com/onnx/models/raw/69c5d3751dda5349fd3fc53f525395d180420c07/vision/classification/inception_and_googlenet/googlenet/model/googlenet-8.onnx"
|
||||
sha1: "da39a3ee5e6b4b0d3255bfef95601890afd80709"
|
||||
model: "googlenet-8.onnx"
|
||||
mean: [103.939, 116.779, 123.675]
|
||||
std: [1, 1, 1]
|
||||
scale: 1.0
|
||||
width: 224
|
||||
height: 224
|
||||
rgb: false
|
||||
labels: "classification_classes_ILSVRC2012.txt"
|
||||
sample: "classification"
|
||||
|
||||
resnet:
|
||||
load_info:
|
||||
url: "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-v2-7.onnx"
|
||||
sha1: "c3a67b3cb2f0a61a7eb75eb8bd9139c89557cbe0"
|
||||
model: "resnet50-v2-7.onnx"
|
||||
mean: [123.675, 116.28, 103.53]
|
||||
std: [58.395, 57.12, 57.375]
|
||||
scale: 1.0
|
||||
width: 224
|
||||
height: 224
|
||||
rgb: true
|
||||
labels: "classification_classes_ILSVRC2012.txt"
|
||||
sample: "classification"
|
||||
|
||||
################################################################################
|
||||
# Semantic segmentation models.
|
||||
################################################################################
|
||||
|
||||
fcnresnet50:
|
||||
load_info:
|
||||
url: "https://github.com/onnx/models/raw/491ce05590abb7551d7fae43c067c060eeb575a6/validated/vision/object_detection_segmentation/fcn/model/fcn-resnet50-12.onnx"
|
||||
sha1: "1bb0c7e0034038969aecc6251166f1612a139230"
|
||||
model: "fcn-resnet50-12.onnx"
|
||||
mean: [103.5, 116.2, 123.6]
|
||||
scale: 0.019
|
||||
width: 500
|
||||
height: 500
|
||||
rgb: false
|
||||
sample: "segmentation"
|
||||
|
||||
fcnresnet101:
|
||||
load_info:
|
||||
url: "https://github.com/onnx/models/raw/fb8271d5d5d9b90dbb1eb5e8e40f8f580fb248b3/vision/object_detection_segmentation/fcn/model/fcn-resnet101-11.onnx"
|
||||
sha1: "e7e76474bf6b73334ab32c4be1374c9e605f5aed"
|
||||
model: "fcn-resnet101-11.onnx"
|
||||
mean: [103.5, 116.2, 123.6]
|
||||
scale: 0.019
|
||||
width: 500
|
||||
height: 500
|
||||
rgb: false
|
||||
sample: "segmentation"
|
||||
|
||||
u2netp:
|
||||
load_info:
|
||||
url: "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
|
||||
sha1: "0a99236f0d5c1916a99a8c401b23e5ef32038606"
|
||||
model: "u2netp.onnx"
|
||||
mean: [123.6, 116.2, 103.5]
|
||||
scale: 0.019
|
||||
width: 320
|
||||
height: 320
|
||||
rgb: true
|
||||
sample: "segmentation"
|
||||
|
||||
################################################################################
|
||||
# Text detection models.
|
||||
################################################################################
|
||||
|
||||
DB:
|
||||
load_info:
|
||||
url: "https://drive.google.com/uc?export=dowload&id=17_ABp79PlFt9yPCxSaarVc_DKTmrSGGf"
|
||||
sha1: "bef233c28947ef6ec8c663d20a2b326302421fa3"
|
||||
model: "DB_IC15_resnet50.onnx"
|
||||
ocr_load_info:
|
||||
ocr_url: "https://drive.google.com/uc?export=dowload&id=159VavnbvfBQkLIPSAu2SP5Yij1Fy4azw"
|
||||
ocr_sha1: "c4ab1fb3f13c1c8ffc04f016e72ec85311de4ebe"
|
||||
ocr_model: "VGG_CTC.onnx"
|
||||
mean: [122.67891434, 116.66876762, 104.00698793]
|
||||
scale: 0.00392
|
||||
width: 736
|
||||
height: 736
|
||||
rgb: false
|
||||
sample: "text_detection"
|
||||
|
||||
East:
|
||||
load_info:
|
||||
url: "https://www.dropbox.com/s/r2ingd0l3zt8hxs/frozen_east_text_detection.tar.gz?dl=1"
|
||||
sha1: "fffabf5ac36f37bddf68e34e84b45f5c4247ed06"
|
||||
download_name: "frozen_east_text_detection.tar.gz"
|
||||
download_sha: "3ca8233d6edd748f7ed23246c8ca24cbf696bb94"
|
||||
model: "frozen_east_text_detection.pb"
|
||||
ocr_load_info:
|
||||
ocr_url: "https://drive.google.com/uc?export=dowload&id=159VavnbvfBQkLIPSAu2SP5Yij1Fy4azw"
|
||||
ocr_sha1: "c4ab1fb3f13c1c8ffc04f016e72ec85311de4ebe"
|
||||
ocr_model: "VGG_CTC.onnx"
|
||||
mean: [123.68, 116.78, 103.94]
|
||||
scale: 1.0
|
||||
width: 736
|
||||
height: 736
|
||||
rgb: false
|
||||
sample: "text_detection"
|
||||
|
||||
OCR:
|
||||
load_info:
|
||||
url: "https://drive.google.com/uc?export=dowload&id=159VavnbvfBQkLIPSAu2SP5Yij1Fy4azw"
|
||||
sha1: "c4ab1fb3f13c1c8ffc04f016e72ec85311de4ebe"
|
||||
model: "VGG_CTC.onnx"
|
||||
sample: "text_recognition"
|
||||
|
||||
# Edge Detection models.
|
||||
################################################################################
|
||||
|
||||
dexined:
|
||||
load_info:
|
||||
url: "https://github.com/opencv/opencv_zoo/raw/refs/heads/main/models/edge_detection_dexined/edge_detection_dexined_2024sep.onnx?download="
|
||||
sha1: "f86f2d32c3cf892771f76b5e6b629b16a66510e9"
|
||||
model: "edge_detection_dexined_2024sep.onnx"
|
||||
mean: [103.5, 116.2, 123.6]
|
||||
scale: 1.0
|
||||
width: 512
|
||||
height: 512
|
||||
rgb: false
|
||||
sample: "edge_detection"
|
||||
|
||||
################################################################################
|
||||
# Edge Detection models.
|
||||
################################################################################
|
||||
|
||||
reid:
|
||||
load_info:
|
||||
url: "https://github.com/opencv/opencv_zoo/raw/main/models/person_reid_youtureid/person_reid_youtu_2021nov.onnx?download="
|
||||
sha1: "d4316b100db40f8840aa82626e1cf3f519a7f1ae"
|
||||
model: "person_reid_youtu_2021nov.onnx"
|
||||
yolo_load_info:
|
||||
yolo_url: "https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov8n.onnx"
|
||||
yolo_sha1: "68f864475d06e2ec4037181052739f268eeac38d"
|
||||
yolo_model: "yolov8n.onnx"
|
||||
mean: [0.485, 0.456, 0.406]
|
||||
std: [0.229, 0.224, 0.225]
|
||||
scale: 0.00392
|
||||
yolo_scale: 0.00392
|
||||
yolo_width: 640
|
||||
yolo_height: 640
|
||||
width: 128
|
||||
height: 256
|
||||
rgb: false
|
||||
yolo_rgb: true
|
||||
sample: "person_reid"
|
||||
|
||||
################################################################################
|
||||
# Tracker models.
|
||||
################################################################################
|
||||
|
||||
vit:
|
||||
load_info:
|
||||
url: "https://github.com/opencv/opencv_zoo/raw/fef72f8fa7c52eaf116d3df358d24e6e959ada0e/models/object_tracking_vittrack/object_tracking_vittrack_2023sep.onnx"
|
||||
sha1: "50008bb4f6a27b1aa940ad886b1bd1936ac4ed3e"
|
||||
model: "object_tracking_vittrack_2023sep.onnx"
|
||||
sample: "object_tracker"
|
||||
|
||||
nanotrack:
|
||||
nanotrack_head_load_info:
|
||||
nanotrack_head_url: "https://github.com/HonglinChu/SiamTrackers/raw/refs/heads/master/NanoTrack/models/nanotrackv2/nanotrack_head_sim.onnx"
|
||||
nanotrack_head_sha1: "39f168489671700cf739e402dfc67d41ce648aef"
|
||||
nanotrack_head_model: "nanotrack_head_sim.onnx"
|
||||
nanotrack_back_load_info:
|
||||
nanotrack_back_url: "https://github.com/HonglinChu/SiamTrackers/raw/refs/heads/master/NanoTrack/models/nanotrackv2/nanotrack_backbone_sim.onnx"
|
||||
nanotrack_back_sha1: "6e773a364457b78574f9f63a23b0659ee8646f8f"
|
||||
nanotrack_back_model: "nanotrack_backbone_sim.onnx"
|
||||
sample: "object_tracker"
|
||||
|
||||
dasiamrpn:
|
||||
dasiamrpn_load_info:
|
||||
dasiamrpn_url: "https://github.com/opencv/opencv_zoo/raw/fef72f8fa7c52eaf116d3df358d24e6e959ada0e/models/object_tracking_dasiamrpn/object_tracking_dasiamrpn_model_2021nov.onnx?download="
|
||||
dasiamrpn_sha1: "91b774fce7df4c0e4918469f0f482d9a27d0e2d4"
|
||||
dasiamrpn_model: "object_tracking_dasiamrpn_model_2021nov.onnx"
|
||||
dasiamrpn_kernel_r1_load_info:
|
||||
dasiamrpn_kernel_r1_url: "https://github.com/opencv/opencv_zoo/raw/fef72f8fa7c52eaf116d3df358d24e6e959ada0e/models/object_tracking_dasiamrpn/object_tracking_dasiamrpn_kernel_r1_2021nov.onnx?download="
|
||||
dasiamrpn_kernel_r1_sha1: "bb64620a54348657133eb28be2d3a2a8c76b84b3"
|
||||
dasiamrpn_kernel_r1_model: "object_tracking_dasiamrpn_kernel_cls1_2021nov.onnx"
|
||||
dasiamrpn_kernel_cls_load_info:
|
||||
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"
|
||||
|
||||
################################################################################
|
||||
# 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"
|
||||
|
||||
ldm_inpainting:
|
||||
encoder_load_info:
|
||||
encoder_url: "https://dl.opencv.org/models/ldm_inpainting/InpaintEncoder.onnx"
|
||||
encoder_sha1: "eb663262304473d81d6ae627d7117892dac56b5e"
|
||||
encoder_model: "InpaintEncoder.onnx"
|
||||
decoder_load_info:
|
||||
decoder_url: "https://dl.opencv.org/models/ldm_inpainting/InpaintDecoder.onnx"
|
||||
decoder_sha1: "af258c100e3a3b0970493b6375c8775beaffc9d1"
|
||||
decoder_model: "InpaintDecoder.onnx"
|
||||
diffusor_load_info:
|
||||
diffusor_url: "https://dl.opencv.org/models/ldm_inpainting/LatentDiffusion.onnx"
|
||||
diffusor_sha1: "2c6f8a505d9a93195510c854d8f023fab27ce70e"
|
||||
diffusor_model: "LatentDiffusion.onnx"
|
||||
mean: [0, 0, 0]
|
||||
scale: 0.00392
|
||||
width: 512
|
||||
height: 512
|
||||
rgb: true
|
||||
sample: "ldm_inpainting"
|
||||
|
||||
################################################################################
|
||||
# Macbeth chart detection model.
|
||||
################################################################################
|
||||
|
||||
mcc:
|
||||
load_info:
|
||||
url: "https://github.com/gursimarsingh/opencv_zoo/raw/refs/heads/mcc_model/models/macbeth_chart_detector/frozen_inference_graph.pb?download="
|
||||
sha1: "fae7dbef14c4ae1fca76f3662220fbd460ed5ed6"
|
||||
model: "frozen_inference_graph.pb"
|
||||
config_load_info:
|
||||
url: "https://github.com/gursimarsingh/opencv_zoo/raw/refs/heads/mcc_model/models/macbeth_chart_detector/graph.pbtxt?download="
|
||||
sha1: "8350cb8f078ecefa1cd566e89930ede25a192310"
|
||||
config: "graph.pbtxt"
|
||||
sample: "mcc"
|
||||
|
||||
################################################################################
|
||||
# Deblurring model.
|
||||
################################################################################
|
||||
|
||||
NAFNet:
|
||||
load_info:
|
||||
url: "https://drive.google.com/uc?export=dowload&id=1ZLRhkpCekNruJZggVpBgSoCx3k7bJ-5v"
|
||||
sha1: "7dabf3d4ede0770ef326afc4511f7e67a791286d"
|
||||
model: "deblurring_nafnet_2025may.onnx"
|
||||
mean: [0, 0, 0]
|
||||
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"
|
||||
|
||||
################################################################################
|
||||
# Super-resolution models.
|
||||
################################################################################
|
||||
|
||||
seemoredetails:
|
||||
load_info:
|
||||
url: "https://github.com/Naresh-19/opencv-superres-models/raw/main/seemore_x4v2_static512.onnx"
|
||||
sha1: "584467bd36f5715aa12c3203bba16e4de6392034"
|
||||
model: "seemore_x4v2_static512.onnx"
|
||||
mean: [0.0, 0.0, 0.0]
|
||||
scale: 0.00392
|
||||
rgb: true
|
||||
width: 512
|
||||
height: 512
|
||||
sample: "super_resolution"
|
||||
input: true
|
||||
|
||||
################################################################################
|
||||
# Auto white balance models.
|
||||
################################################################################
|
||||
|
||||
fc4:
|
||||
load_info:
|
||||
url: "https://raw.githubusercontent.com/MykhailoTrushch/fc4-models/main/fc4_fold_0.onnx"
|
||||
sha1: "e8a9a65ec0baaae3e4c97b34274a620eb362e905"
|
||||
model: "fc4_fold_0.onnx"
|
||||
sample: "auto_white_balance"
|
||||
scale: 0.00392156862
|
||||
rgb: true
|
||||
mean: 0
|
||||
@@ -0,0 +1,735 @@
|
||||
//![includes]
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/core/utils/logger.hpp>
|
||||
|
||||
#include <mutex>
|
||||
#include <thread>
|
||||
#include <queue>
|
||||
|
||||
#include "iostream"
|
||||
#include "common.hpp"
|
||||
//![includes]
|
||||
|
||||
using namespace cv;
|
||||
using namespace dnn;
|
||||
using namespace std;
|
||||
|
||||
const string about =
|
||||
"Firstly, download required models 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.\n"
|
||||
"To run:\n"
|
||||
"\t ./example_dnn_object_detection model_name --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)\n"
|
||||
"Sample command:\n"
|
||||
"\t ./example_dnn_object_detection yolov8 --input=$OPENCV_SAMPLES_DATA_PATH/baboon.jpg\n"
|
||||
|
||||
"Model path can also be specified using --model argument. ";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ @alias | | 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 }"
|
||||
"{ device | 0 | camera device number. }"
|
||||
"{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }"
|
||||
"{ thr | .5 | Confidence threshold. }"
|
||||
"{ nms | .4 | Non-maximum suppression threshold. }"
|
||||
"{ async | 0 | Number of asynchronous forwards at the same time. "
|
||||
"Choose 0 for synchronous mode }"
|
||||
"{ padvalue | 114.0 | padding value. }"
|
||||
"{ paddingmode | 2 | Choose one of padding modes: "
|
||||
"0: resize to required input size without extra processing, "
|
||||
"1: Image will be cropped after resize, "
|
||||
"2: Resize image to the desired size while preserving the aspect ratio of original image }";
|
||||
|
||||
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;
|
||||
|
||||
float confThreshold, nmsThreshold, scale, paddingValue;
|
||||
vector<string> labels;
|
||||
Scalar meanv;
|
||||
bool swapRB;
|
||||
int inpWidth, inpHeight;
|
||||
size_t asyncNumReq = 0;
|
||||
ImagePaddingMode paddingMode;
|
||||
string modelName, framework;
|
||||
|
||||
static void preprocess(const Mat& frame, Net& net, Size inpSize);
|
||||
|
||||
static void postprocess(Mat& frame, const vector<Mat>& outs, Net& net, vector<int>& classIds, vector<float>& confidences, vector<Rect>& boxes, const string postprocessing);
|
||||
|
||||
static void drawPred(vector<int>& classIds, vector<float>& confidences, vector<Rect>& boxes, Mat& frame, FontFace& sans, int stdSize, int stdWeight, int stdImgSize, int stdThickness);
|
||||
|
||||
static void callback(int pos, void* userdata);
|
||||
|
||||
static Scalar getColor(int classId);
|
||||
|
||||
static void yoloPostProcessing(
|
||||
const vector<Mat>& outs,
|
||||
vector<int>& keep_classIds,
|
||||
vector<float>& keep_confidences,
|
||||
vector<Rect2d>& keep_boxes,
|
||||
float conf_threshold,
|
||||
float iou_threshold,
|
||||
const string& postprocessing);
|
||||
|
||||
static void printAliases(string& zooFile){
|
||||
vector<string> aliases = findAliases(zooFile, "object_detection");
|
||||
|
||||
cout<<"Alias choices: [ ";
|
||||
for (auto it: aliases){
|
||||
cout<<"'"<<it<<"' ";
|
||||
}
|
||||
cout<<"]"<<endl;
|
||||
}
|
||||
|
||||
static Scalar getTextColor(Scalar bgColor) {
|
||||
double luminance = 0.299 * bgColor[2] + 0.587 * bgColor[1] + 0.114 * bgColor[0];
|
||||
|
||||
return luminance > 128 ? Scalar(0, 0, 0) : Scalar(255, 255, 255);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
class QueueFPS : public std::queue<T>
|
||||
{
|
||||
public:
|
||||
QueueFPS() : counter(0) {}
|
||||
|
||||
void push(const T& entry)
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
|
||||
std::queue<T>::push(entry);
|
||||
counter += 1;
|
||||
if (counter == 1)
|
||||
{
|
||||
// Start counting from a second frame (warmup).
|
||||
tm.reset();
|
||||
tm.start();
|
||||
}
|
||||
}
|
||||
|
||||
T get()
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
T entry = this->front();
|
||||
this->pop();
|
||||
return entry;
|
||||
}
|
||||
|
||||
float getFPS()
|
||||
{
|
||||
tm.stop();
|
||||
double fps = counter / tm.getTimeSec();
|
||||
tm.start();
|
||||
return static_cast<float>(fps);
|
||||
}
|
||||
|
||||
void clear()
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
while (!this->empty())
|
||||
this->pop();
|
||||
}
|
||||
|
||||
unsigned int counter;
|
||||
|
||||
private:
|
||||
TickMeter tm;
|
||||
std::mutex mutex;
|
||||
};
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
utils::logging::setLogLevel(utils::logging::LOG_LEVEL_INFO);
|
||||
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
|
||||
string zooFile = parser.get<String>("zoo");
|
||||
if (!parser.has("@alias") || parser.has("help"))
|
||||
{
|
||||
cout << about << endl;
|
||||
parser.printMessage();
|
||||
printAliases(zooFile);
|
||||
return -1;
|
||||
}
|
||||
zooFile = findFile(zooFile);
|
||||
modelName = parser.get<String>("@alias");
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
|
||||
if (!parser.has("model"))
|
||||
{
|
||||
cout << "Path to model is not provided in command line or model alias is not correct" << endl;
|
||||
printAliases(zooFile);
|
||||
return -1;
|
||||
}
|
||||
|
||||
confThreshold = parser.get<float>("thr");
|
||||
nmsThreshold = parser.get<float>("nms");
|
||||
//![preprocess_params]
|
||||
scale = parser.get<float>("scale");
|
||||
meanv = parser.get<Scalar>("mean");
|
||||
swapRB = parser.get<bool>("rgb");
|
||||
inpWidth = parser.get<int>("width");
|
||||
inpHeight = parser.get<int>("height");
|
||||
int async = parser.get<int>("async");
|
||||
paddingValue = parser.get<float>("padvalue");
|
||||
const string postprocessing = parser.get<String>("postprocessing");
|
||||
paddingMode = static_cast<ImagePaddingMode>(parser.get<int>("paddingmode"));
|
||||
//![preprocess_params]
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
String config_sha1 = parser.get<String>("config_sha1");
|
||||
const string modelPath = findModel(parser.get<String>("model"), sha1);
|
||||
const string configPath = findModel(parser.get<String>("config"), config_sha1);
|
||||
framework = modelPath.substr(modelPath.rfind('.') + 1);
|
||||
|
||||
if (parser.has("labels"))
|
||||
{
|
||||
const string file = findFile(parser.get<String>("labels"));
|
||||
ifstream ifs(file.c_str());
|
||||
if (!ifs.is_open())
|
||||
CV_Error(Error::StsError, "File " + file + " not found");
|
||||
string line;
|
||||
while (getline(ifs, line))
|
||||
{
|
||||
labels.push_back(line);
|
||||
}
|
||||
}
|
||||
//![read_net]
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if ((parser.get<String>("backend") != "default") || (parser.get<String>("target") != "cpu")){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
Net net = readNet(modelPath, configPath, "", engine);
|
||||
int backend = getBackendID(parser.get<String>("backend"));
|
||||
net.setPreferableBackend(backend);
|
||||
net.setPreferableTarget(getTargetID(parser.get<String>("target")));
|
||||
net.setProfilingMode(DNN_PROFILE_SUMMARY);
|
||||
//![read_net]
|
||||
|
||||
// Create a window
|
||||
static const string kWinName = "Deep learning object detection in OpenCV";
|
||||
namedWindow(kWinName, WINDOW_AUTOSIZE);
|
||||
int initialConf = (int)(confThreshold * 100);
|
||||
createTrackbar("Confidence threshold, %", kWinName, &initialConf, 99, callback, &net);
|
||||
|
||||
// Open a video file or an image file or a camera stream.
|
||||
VideoCapture cap;
|
||||
bool openSuccess = parser.has("input") ? cap.open(findFile(parser.get<String>("input"))) : cap.open(parser.get<int>("device"));
|
||||
if (!openSuccess){
|
||||
cout << "Could not open input file or camera device" << endl;
|
||||
return 0;
|
||||
}
|
||||
|
||||
FontFace sans("sans");
|
||||
|
||||
int stdSize = 15;
|
||||
int stdWeight = 150;
|
||||
int stdImgSize = 512;
|
||||
int stdThickness = 2;
|
||||
vector<int> classIds;
|
||||
vector<float> confidences;
|
||||
vector<Rect> boxes;
|
||||
|
||||
if (async > 0 && backend == DNN_BACKEND_INFERENCE_ENGINE){
|
||||
asyncNumReq = async;
|
||||
}
|
||||
|
||||
if (async != 0) {
|
||||
// Threading is enabled
|
||||
bool process = true;
|
||||
|
||||
// Frames capturing thread
|
||||
QueueFPS<Mat> framesQueue;
|
||||
std::thread framesThread([&]() {
|
||||
Mat frame;
|
||||
while (process) {
|
||||
cap >> frame;
|
||||
if (!frame.empty())
|
||||
framesQueue.push(frame.clone());
|
||||
else
|
||||
break;
|
||||
}
|
||||
});
|
||||
|
||||
// Frames processing thread
|
||||
QueueFPS<Mat> processedFramesQueue;
|
||||
QueueFPS<std::vector<Mat>> predictionsQueue;
|
||||
std::thread processingThread([&]() {
|
||||
std::queue<AsyncArray> futureOutputs;
|
||||
Mat blob;
|
||||
while (process) {
|
||||
// Get the next frame
|
||||
Mat frame;
|
||||
{
|
||||
if (!framesQueue.empty()) {
|
||||
frame = framesQueue.get();
|
||||
if (asyncNumReq) {
|
||||
if (futureOutputs.size() == asyncNumReq)
|
||||
frame = Mat();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process the frame
|
||||
if (!frame.empty()) {
|
||||
preprocess(frame, net, Size(inpWidth, inpHeight));
|
||||
processedFramesQueue.push(frame);
|
||||
|
||||
if (asyncNumReq) {
|
||||
futureOutputs.push(net.forwardAsync());
|
||||
} else {
|
||||
//![forward]
|
||||
vector<Mat> outs;
|
||||
net.forward(outs, net.getUnconnectedOutLayersNames());
|
||||
net.printPerfProfile();
|
||||
predictionsQueue.push(outs);
|
||||
//![forward]
|
||||
}
|
||||
}
|
||||
|
||||
while (!futureOutputs.empty() &&
|
||||
futureOutputs.front().wait_for(std::chrono::seconds(0))) {
|
||||
AsyncArray async_out = futureOutputs.front();
|
||||
futureOutputs.pop();
|
||||
Mat out;
|
||||
async_out.get(out);
|
||||
predictionsQueue.push({out});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Postprocessing and rendering loop
|
||||
while (waitKey(100) < 0) {
|
||||
if (predictionsQueue.empty())
|
||||
continue;
|
||||
|
||||
vector<Mat> outs = predictionsQueue.get();
|
||||
Mat frame = processedFramesQueue.get();
|
||||
|
||||
classIds.clear();
|
||||
confidences.clear();
|
||||
boxes.clear();
|
||||
postprocess(frame, outs, net, classIds, confidences, boxes, postprocessing);
|
||||
|
||||
drawPred(classIds, confidences, boxes, frame, sans, stdSize, stdWeight, stdImgSize, stdThickness);
|
||||
|
||||
int imgWidth = max(frame.rows, frame.cols);
|
||||
int size = static_cast<int>((stdSize * imgWidth) / (stdImgSize * 1.5));
|
||||
int weight = static_cast<int>((stdWeight * imgWidth) / (stdImgSize * 1.5));
|
||||
|
||||
if (predictionsQueue.counter > 1) {
|
||||
string label = format("Camera: %.2f FPS", framesQueue.getFPS());
|
||||
rectangle(frame, Point(0, 0), Point(10 * size, 3 * size + size / 4), Scalar::all(255), FILLED);
|
||||
putText(frame, label, Point(0, size), Scalar::all(0), sans, size, weight);
|
||||
|
||||
label = format("Network: %.2f FPS", predictionsQueue.getFPS());
|
||||
putText(frame, label, Point(0, 2 * size), Scalar::all(0), sans, size, weight);
|
||||
|
||||
label = format("Skipped frames: %d", framesQueue.counter - predictionsQueue.counter);
|
||||
putText(frame, label, Point(0, 3 * size), Scalar::all(0), sans, size, weight);
|
||||
}
|
||||
imshow(kWinName, frame);
|
||||
}
|
||||
|
||||
process = false;
|
||||
framesThread.join();
|
||||
processingThread.join();
|
||||
} else {
|
||||
if (asyncNumReq)
|
||||
CV_Error(Error::StsNotImplemented, "Asynchronous forward is supported only with Inference Engine backend.");
|
||||
// Threading is disabled, run synchronously
|
||||
Mat frame, blob;
|
||||
while (waitKey(1) < 0) {
|
||||
cap >> frame;
|
||||
if (frame.empty()) {
|
||||
waitKey();
|
||||
break;
|
||||
}
|
||||
preprocess(frame, net, Size(inpWidth, inpHeight));
|
||||
|
||||
TickMeter tickMeter;
|
||||
vector<Mat> outs;
|
||||
tickMeter.start();
|
||||
net.forward(outs, net.getUnconnectedOutLayersNames());
|
||||
tickMeter.stop();
|
||||
net.printPerfProfile();
|
||||
|
||||
classIds.clear();
|
||||
confidences.clear();
|
||||
boxes.clear();
|
||||
|
||||
postprocess(frame, outs, net, classIds, confidences, boxes, postprocessing);
|
||||
|
||||
drawPred(classIds, confidences, boxes, frame, sans, stdSize, stdWeight, stdImgSize, stdThickness);
|
||||
|
||||
int imgWidth = max(frame.rows, frame.cols);
|
||||
int size = static_cast<int>((stdSize * imgWidth) / (stdImgSize * 1.5));
|
||||
int weight = static_cast<int>((stdWeight * imgWidth) / (stdImgSize * 1.5));
|
||||
string label = format("FPS: %.2f", 1000.0 / tickMeter.getTimeMilli());
|
||||
putText(frame, label, Point(0, size), Scalar(0, 255, 0), sans, size, weight);
|
||||
imshow(kWinName, frame);
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
void preprocess(const Mat& frame, Net& net, Size inpSize)
|
||||
{
|
||||
Size size(inpSize.width <= 0 ? frame.cols : inpSize.width, inpSize.height <= 0 ? frame.rows : inpSize.height);
|
||||
|
||||
// Prepare the blob from the image
|
||||
Mat inp;
|
||||
{
|
||||
//![preprocess_call]
|
||||
Image2BlobParams imgParams(
|
||||
Scalar::all(scale),
|
||||
size,
|
||||
meanv,
|
||||
swapRB,
|
||||
CV_32F,
|
||||
DNN_LAYOUT_NCHW,
|
||||
paddingMode,
|
||||
paddingValue);
|
||||
|
||||
inp = blobFromImageWithParams(frame, imgParams);
|
||||
//![preprocess_call]
|
||||
}
|
||||
|
||||
// Set the blob as the network input
|
||||
net.setInput(inp);
|
||||
}
|
||||
|
||||
void yoloPostProcessing(
|
||||
const vector<Mat>& outs,
|
||||
vector<int>& keep_classIds,
|
||||
vector<float>& keep_confidences,
|
||||
vector<Rect2d>& keep_boxes,
|
||||
float conf_threshold,
|
||||
float iou_threshold,
|
||||
const string& postprocessing)
|
||||
{
|
||||
// Retrieve
|
||||
vector<int> classIds;
|
||||
vector<float> confidences;
|
||||
vector<Rect2d> boxes;
|
||||
|
||||
vector<Mat> outs_copy = outs;
|
||||
|
||||
if (postprocessing == "yolov8")
|
||||
{
|
||||
transposeND(outs_copy[0], {0, 2, 1}, outs_copy[0]);
|
||||
}
|
||||
|
||||
if (postprocessing == "yolonas")
|
||||
{
|
||||
// outs contains 2 elements of shape [1, 8400, 80] and [1, 8400, 4]. Concat them to get [1, 8400, 84]
|
||||
Mat concat_out;
|
||||
// squeeze the first dimension
|
||||
outs_copy[0] = outs_copy[0].reshape(1, outs_copy[0].size[1]);
|
||||
outs_copy[1] = outs_copy[1].reshape(1, outs_copy[1].size[1]);
|
||||
hconcat(outs_copy[1], outs_copy[0], concat_out);
|
||||
outs_copy[0] = concat_out;
|
||||
// remove the second element
|
||||
outs_copy.pop_back();
|
||||
// unsqueeze the first dimension
|
||||
outs_copy[0] = outs_copy[0].reshape(0, vector<int>{1, 8400, 84});
|
||||
}
|
||||
|
||||
for (auto preds : outs_copy)
|
||||
{
|
||||
preds = preds.reshape(1, preds.size[1]); // [1, 8400, 85] -> [8400, 85]
|
||||
for (int i = 0; i < preds.rows; ++i)
|
||||
{
|
||||
// filter out non-object
|
||||
float obj_conf = (postprocessing == "yolov8" || postprocessing == "yolonas") ? 1.0f : preds.at<float>(i, 4);
|
||||
if (obj_conf < conf_threshold)
|
||||
continue;
|
||||
|
||||
Mat scores = preds.row(i).colRange((postprocessing == "yolov8" || postprocessing == "yolonas") ? 4 : 5, preds.cols);
|
||||
double conf;
|
||||
Point maxLoc;
|
||||
minMaxLoc(scores, 0, &conf, 0, &maxLoc);
|
||||
|
||||
conf = (postprocessing == "yolov8" || postprocessing == "yolonas") ? conf : conf * obj_conf;
|
||||
if (conf < conf_threshold)
|
||||
continue;
|
||||
|
||||
// get bbox coords
|
||||
float* det = preds.ptr<float>(i);
|
||||
double cx = det[0];
|
||||
double cy = det[1];
|
||||
double w = det[2];
|
||||
double h = det[3];
|
||||
|
||||
// [x1, y1, x2, y2]
|
||||
if (postprocessing == "yolonas") {
|
||||
boxes.push_back(Rect2d(cx, cy, w, h));
|
||||
} else {
|
||||
boxes.push_back(Rect2d(cx - 0.5 * w, cy - 0.5 * h,
|
||||
cx + 0.5 * w, cy + 0.5 * h));
|
||||
}
|
||||
classIds.push_back(maxLoc.x);
|
||||
confidences.push_back(static_cast<float>(conf));
|
||||
}
|
||||
}
|
||||
|
||||
// NMS
|
||||
vector<int> keep_idx;
|
||||
NMSBoxes(boxes, confidences, conf_threshold, iou_threshold, keep_idx);
|
||||
|
||||
for (auto i : keep_idx)
|
||||
{
|
||||
keep_classIds.push_back(classIds[i]);
|
||||
keep_confidences.push_back(confidences[i]);
|
||||
keep_boxes.push_back(boxes[i]);
|
||||
}
|
||||
}
|
||||
|
||||
void postprocess(Mat& frame, const vector<Mat>& outs, Net& net, vector<int>& classIds, vector<float>& confidences, vector<Rect>& boxes, const string postprocessing)
|
||||
{
|
||||
static vector<int> outLayers = net.getUnconnectedOutLayers();
|
||||
if (postprocessing == "ssd")
|
||||
{
|
||||
// Network produces output blob with a shape 1x1xNx7 where N is a number of
|
||||
// detections and an every detection is a vector of values
|
||||
// [batchId, classId, confidence, left, top, right, bottom]
|
||||
CV_Assert(outs.size() > 0);
|
||||
for (size_t k = 0; k < outs.size(); k++)
|
||||
{
|
||||
float* data = (float*)outs[k].data;
|
||||
for (size_t i = 0; i < outs[k].total(); i += 7)
|
||||
{
|
||||
float confidence = data[i + 2];
|
||||
if (confidence > confThreshold)
|
||||
{
|
||||
int left = (int)data[i + 3];
|
||||
int top = (int)data[i + 4];
|
||||
int right = (int)data[i + 5];
|
||||
int bottom = (int)data[i + 6];
|
||||
int width = right - left + 1;
|
||||
int height = bottom - top + 1;
|
||||
if (width <= 2 || height <= 2)
|
||||
{
|
||||
left = (int)(data[i + 3] * frame.cols);
|
||||
top = (int)(data[i + 4] * frame.rows);
|
||||
right = (int)(data[i + 5] * frame.cols);
|
||||
bottom = (int)(data[i + 6] * frame.rows);
|
||||
width = right - left + 1;
|
||||
height = bottom - top + 1;
|
||||
}
|
||||
classIds.push_back((int)(data[i + 1]) - 1); // Skip 0th background class id.
|
||||
boxes.push_back(Rect(left, top, width, height));
|
||||
confidences.push_back(confidence);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
else if (postprocessing == "yolov4")
|
||||
{
|
||||
// boxes[b,N,1,4]+confs[b,N,classes] (normalized) or boxes[b,N,4]+scores[b,N]+classIdx[b,N] (model-px)
|
||||
bool isBoxConfsFormat = (outs.size() == 2 && outs[0].dims == 4 && outs[0].size[outs[0].dims - 1] == 4);
|
||||
bool isBoxScoresIdxFormat = (outs.size() == 3 && outs[0].dims == 3 && outs[0].size[2] == 4);
|
||||
if (isBoxScoresIdxFormat)
|
||||
{
|
||||
int N = outs[0].size[1];
|
||||
const float* boxesPtr = outs[0].ptr<float>(0);
|
||||
const float* scoresPtr = outs[1].ptr<float>(0);
|
||||
const float* classIdxPtr = outs[2].ptr<float>(0);
|
||||
for (int j = 0; j < N; ++j)
|
||||
{
|
||||
float score = scoresPtr[j];
|
||||
if (score > confThreshold)
|
||||
{
|
||||
float x1 = boxesPtr[j * 4 + 0];
|
||||
float y1 = boxesPtr[j * 4 + 1];
|
||||
float x2 = boxesPtr[j * 4 + 2];
|
||||
float y2 = boxesPtr[j * 4 + 3];
|
||||
boxes.push_back(Rect((int)x1, (int)y1, (int)(x2 - x1), (int)(y2 - y1)));
|
||||
confidences.push_back(score);
|
||||
classIds.push_back((int)classIdxPtr[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
else if (isBoxConfsFormat)
|
||||
{
|
||||
Mat boxesMat = outs[0];
|
||||
Mat confsMat = outs[1];
|
||||
int numBoxes = (int)(boxesMat.total() / 4);
|
||||
boxesMat = boxesMat.reshape(1, numBoxes);
|
||||
confsMat = confsMat.reshape(1, numBoxes);
|
||||
for (int j = 0; j < numBoxes; ++j)
|
||||
{
|
||||
Point maxLoc;
|
||||
double confidence;
|
||||
minMaxLoc(confsMat.row(j), 0, &confidence, 0, &maxLoc);
|
||||
if (confidence > confThreshold)
|
||||
{
|
||||
const float* box = boxesMat.ptr<float>(j);
|
||||
boxes.push_back(Rect((int)(box[0] * inpWidth), (int)(box[1] * inpHeight),
|
||||
(int)((box[2] - box[0]) * inpWidth), (int)((box[3] - box[1]) * inpHeight)));
|
||||
confidences.push_back((float)confidence);
|
||||
classIds.push_back(maxLoc.x);
|
||||
}
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
cout << "Unsupported YOLO ONNX output format" << endl;
|
||||
exit(-1);
|
||||
}
|
||||
Image2BlobParams paramNet;
|
||||
paramNet.scalefactor = Scalar::all(scale);
|
||||
paramNet.size = Size(inpWidth, inpHeight);
|
||||
paramNet.mean = meanv;
|
||||
paramNet.swapRB = swapRB;
|
||||
paramNet.paddingmode = paddingMode;
|
||||
paramNet.blobRectsToImageRects(boxes, boxes, frame.size());
|
||||
}
|
||||
else if (postprocessing == "yolov8" || postprocessing == "yolov5")
|
||||
{
|
||||
//![forward_buffers]
|
||||
vector<int> keep_classIds;
|
||||
vector<float> keep_confidences;
|
||||
vector<Rect2d> keep_boxes;
|
||||
//![forward_buffers]
|
||||
|
||||
//![postprocess]
|
||||
yoloPostProcessing(outs, keep_classIds, keep_confidences, keep_boxes, confThreshold, nmsThreshold, postprocessing);
|
||||
//![postprocess]
|
||||
|
||||
for (size_t i = 0; i < keep_classIds.size(); ++i)
|
||||
{
|
||||
classIds.push_back(keep_classIds[i]);
|
||||
confidences.push_back(keep_confidences[i]);
|
||||
Rect2d box = keep_boxes[i];
|
||||
boxes.push_back(Rect(cvFloor(box.x), cvFloor(box.y), cvFloor(box.width-box.x), cvFloor(box.height-box.y)));
|
||||
}
|
||||
if (framework == "onnx"){
|
||||
Image2BlobParams paramNet;
|
||||
paramNet.scalefactor = Scalar::all(scale);
|
||||
paramNet.size = Size(inpWidth, inpHeight);
|
||||
paramNet.mean = meanv;
|
||||
paramNet.swapRB = swapRB;
|
||||
paramNet.paddingmode = paddingMode;
|
||||
|
||||
paramNet.blobRectsToImageRects(boxes, boxes, frame.size());
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
cout<< ("Unknown postprocessing method: " + postprocessing)<<endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
// NMS is used inside Region layer only on DNN_BACKEND_OPENCV for other backends we need NMS in sample
|
||||
// or NMS is required if the number of outputs > 1
|
||||
if (outLayers.size() > 1)
|
||||
{
|
||||
map<int, vector<size_t> > class2indices;
|
||||
for (size_t i = 0; i < classIds.size(); i++)
|
||||
{
|
||||
if (confidences[i] >= confThreshold)
|
||||
{
|
||||
class2indices[classIds[i]].push_back(i);
|
||||
}
|
||||
}
|
||||
vector<Rect> nmsBoxes;
|
||||
vector<float> nmsConfidences;
|
||||
vector<int> nmsClassIds;
|
||||
for (map<int, vector<size_t> >::iterator it = class2indices.begin(); it != class2indices.end(); ++it)
|
||||
{
|
||||
vector<Rect> localBoxes;
|
||||
vector<float> localConfidences;
|
||||
vector<size_t> classIndices = it->second;
|
||||
for (size_t i = 0; i < classIndices.size(); i++)
|
||||
{
|
||||
localBoxes.push_back(boxes[classIndices[i]]);
|
||||
localConfidences.push_back(confidences[classIndices[i]]);
|
||||
}
|
||||
vector<int> nmsIndices;
|
||||
NMSBoxes(localBoxes, localConfidences, confThreshold, nmsThreshold, nmsIndices);
|
||||
for (size_t i = 0; i < nmsIndices.size(); i++)
|
||||
{
|
||||
size_t idx = nmsIndices[i];
|
||||
nmsBoxes.push_back(localBoxes[idx]);
|
||||
nmsConfidences.push_back(localConfidences[idx]);
|
||||
nmsClassIds.push_back(it->first);
|
||||
}
|
||||
}
|
||||
boxes = nmsBoxes;
|
||||
classIds = nmsClassIds;
|
||||
confidences = nmsConfidences;
|
||||
}
|
||||
}
|
||||
|
||||
void drawPred(vector<int>& classIds, vector<float>& confidences, vector<Rect>& boxes, Mat& frame, FontFace& sans, int stdSize, int stdWeight, int stdImgSize, int stdThickness)
|
||||
{
|
||||
//![draw_boxes]
|
||||
int imgWidth = max(frame.rows, frame.cols);
|
||||
int size = (stdSize*imgWidth)/stdImgSize;
|
||||
int weight = (stdWeight*imgWidth)/stdImgSize;
|
||||
int thickness = (stdThickness*imgWidth)/stdImgSize;
|
||||
|
||||
for (size_t idx = 0; idx < boxes.size(); ++idx){
|
||||
Scalar boxColor = getColor(classIds[idx]);
|
||||
int left = boxes[idx].x;
|
||||
int top = boxes[idx].y;
|
||||
int right = boxes[idx].x + boxes[idx].width;
|
||||
int bottom = boxes[idx].y + boxes[idx].height;
|
||||
rectangle(frame, Point(left, top), Point(right, bottom), boxColor, thickness);
|
||||
|
||||
string label = format("%.2f", confidences[idx]);
|
||||
if (!labels.empty())
|
||||
{
|
||||
CV_Assert(classIds[idx] < (int)labels.size());
|
||||
label = labels[classIds[idx]] + ": " + label;
|
||||
}
|
||||
|
||||
Rect r = getTextSize(Size(), label, Point(), sans, size, weight);
|
||||
int baseline = r.y + r.height;
|
||||
Size labelSize = Size(r.width, r.height + size/4 - baseline);
|
||||
|
||||
top = max(top-thickness/2, labelSize.height);
|
||||
rectangle(frame, Point(left-thickness/2, top-(labelSize.height)),
|
||||
Point(left + labelSize.width, top), boxColor, FILLED);
|
||||
putText(frame, label, Point(left, top-size/4), getTextColor(boxColor), sans, size, weight);
|
||||
}
|
||||
//![draw_boxes]
|
||||
}
|
||||
|
||||
void callback(int pos, void*)
|
||||
{
|
||||
confThreshold = pos * 0.01f;
|
||||
}
|
||||
|
||||
Scalar getColor(int classId) {
|
||||
int r = min((classId >> 0 & 1) * 128 + (classId >> 3 & 1) * 64 + (classId >> 6 & 1) * 32 + 80, 255);
|
||||
int g = min((classId >> 1 & 1) * 128 + (classId >> 4 & 1) * 64 + (classId >> 7 & 1) * 32 + 40, 255);
|
||||
int b = min((classId >> 2 & 1) * 128 + (classId >> 5 & 1) * 64 + (classId >> 8 & 1) * 32 + 40, 255);
|
||||
return Scalar(b, g, r);
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import numpy as np
|
||||
import sys
|
||||
import copy
|
||||
import time
|
||||
from threading import Thread
|
||||
import queue
|
||||
|
||||
from common import *
|
||||
from tf_text_graph_common import readTextMessage
|
||||
from tf_text_graph_ssd import createSSDGraph
|
||||
from tf_text_graph_faster_rcnn import createFasterRCNNGraph
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Firstly, download required models 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.\n"\n
|
||||
|
||||
To run:
|
||||
python object_detection.py model_name(e.g yolov8) --input=path/to/your/input/image/or/video (don't pass --input to use device camera)
|
||||
|
||||
Sample command:
|
||||
python object_detection.py yolov8 --input=path/to/image
|
||||
Model path can also be specified using --model argument
|
||||
'''
|
||||
)
|
||||
|
||||
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', help='Path to input image or video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--out_tf_graph', default='graph.pbtxt',
|
||||
help='For models from TensorFlow Object Detection API, you may '
|
||||
'pass a .config file which was used for training through --config '
|
||||
'argument. This way an additional .pbtxt file with TensorFlow graph will be created.')
|
||||
parser.add_argument('--thr', type=float, default=0.5, help='Confidence threshold')
|
||||
parser.add_argument('--nms', type=float, default=0.4, help='Non-maximum suppression threshold')
|
||||
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)")
|
||||
parser.add_argument('--async', type=int, default=0,
|
||||
dest='use_threads',
|
||||
help='Choose 0 for synchronous mode and 1 for asynchronous mode')
|
||||
args, _ = parser.parse_known_args()
|
||||
add_preproc_args(args.zoo, parser, 'object_detection')
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Use this script to run object detection deep learning networks using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.alias is None or hasattr(args, 'help'):
|
||||
help()
|
||||
exit(1)
|
||||
|
||||
cv.utils.logging.setLogLevel(cv.utils.logging.LOG_LEVEL_INFO)
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
if args.config is not None:
|
||||
args.config = findModel(args.config, args.config_sha1)
|
||||
if args.labels is not None:
|
||||
args.labels = findFile(args.labels)
|
||||
|
||||
# If config specified, try to load it as TensorFlow Object Detection API's pipeline.
|
||||
config = readTextMessage(args.config)
|
||||
if 'model' in config:
|
||||
print('TensorFlow Object Detection API config detected')
|
||||
if 'ssd' in config['model'][0]:
|
||||
print('Preparing text graph representation for SSD model: ' + args.out_tf_graph)
|
||||
createSSDGraph(args.model, args.config, args.out_tf_graph)
|
||||
args.config = args.out_tf_graph
|
||||
elif 'faster_rcnn' in config['model'][0]:
|
||||
print('Preparing text graph representation for Faster-RCNN model: ' + args.out_tf_graph)
|
||||
createFasterRCNNGraph(args.model, args.config, args.out_tf_graph)
|
||||
args.config = args.out_tf_graph
|
||||
|
||||
|
||||
# Load names of classes
|
||||
labels = None
|
||||
if args.labels:
|
||||
with open(args.labels, 'rt') as f:
|
||||
labels = f.read().rstrip('\n').split('\n')
|
||||
|
||||
# Load a network
|
||||
engine = cv.dnn.ENGINE_AUTO
|
||||
if args.backend != "default" or args.target != "cpu":
|
||||
engine = cv.dnn.ENGINE_CLASSIC
|
||||
net = cv.dnn.readNet(args.model, args.config, "", engine)
|
||||
net.setPreferableBackend(get_backend_id(args.backend))
|
||||
net.setPreferableTarget(get_target_id(args.target))
|
||||
if hasattr(cv.dnn, 'DNN_PROFILE_SUMMARY'):
|
||||
net.setProfilingMode(cv.dnn.DNN_PROFILE_SUMMARY)
|
||||
outNames = net.getUnconnectedOutLayersNames()
|
||||
|
||||
confThreshold = args.thr
|
||||
nmsThreshold = args.nms
|
||||
stdSize = 0.8
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
asyncN = 0
|
||||
|
||||
def get_color(class_id):
|
||||
r = min((class_id >> 0 & 1) * 128 + (class_id >> 3 & 1) * 64 + (class_id >> 6 & 1) * 32 + 80, 255)
|
||||
g = min((class_id >> 1 & 1) * 128 + (class_id >> 4 & 1) * 64 + (class_id >> 7 & 1) * 32 + 40, 255)
|
||||
b = min((class_id >> 2 & 1) * 128 + (class_id >> 5 & 1) * 64 + (class_id >> 8 & 1) * 32 + 40, 255)
|
||||
return (int(b), int(g), int(r))
|
||||
|
||||
def get_text_color(bg_color):
|
||||
luminance = 0.299 * bg_color[2] + 0.587 * bg_color[1] + 0.114 * bg_color[0]
|
||||
return (0, 0, 0) if luminance > 128 else (255, 255, 255)
|
||||
|
||||
def postprocess(frame, outs):
|
||||
frameHeight = frame.shape[0]
|
||||
frameWidth = frame.shape[1]
|
||||
|
||||
classIds = []
|
||||
confidences = []
|
||||
boxes = []
|
||||
if args.postprocessing == 'ssd':
|
||||
# Network produces output blob with a shape 1x1xNx7 where N is a number of
|
||||
# detections and an every detection is a vector of values
|
||||
# [batchId, classId, confidence, left, top, right, bottom]
|
||||
for out in outs:
|
||||
for detection in out[0, 0]:
|
||||
confidence = detection[2]
|
||||
if confidence > confThreshold:
|
||||
left = int(detection[3])
|
||||
top = int(detection[4])
|
||||
right = int(detection[5])
|
||||
bottom = int(detection[6])
|
||||
width = right - left + 1
|
||||
height = bottom - top + 1
|
||||
if width <= 2 or height <= 2:
|
||||
left = int(detection[3] * frameWidth)
|
||||
top = int(detection[4] * frameHeight)
|
||||
right = int(detection[5] * frameWidth)
|
||||
bottom = int(detection[6] * frameHeight)
|
||||
width = right - left + 1
|
||||
height = bottom - top + 1
|
||||
classIds.append(int(detection[1]) - 1) # Skip background label
|
||||
confidences.append(float(confidence))
|
||||
boxes.append([left, top, width, height])
|
||||
|
||||
elif args.postprocessing == 'yolov4':
|
||||
# boxes[b,N,1,4]+confs[b,N,classes] (normalized) or boxes[b,N,4]+scores[b,N]+classIdx[b,N] (model-px)
|
||||
if len(outs) == 3 and outs[0].ndim == 3 and outs[0].shape[2] == 4:
|
||||
boxesArr = outs[0][0]
|
||||
scoresArr = outs[1][0]
|
||||
classIdxArr = outs[2][0]
|
||||
for j in range(boxesArr.shape[0]):
|
||||
score = float(scoresArr[j])
|
||||
if score > confThreshold:
|
||||
x1 = boxesArr[j][0] / args.width
|
||||
y1 = boxesArr[j][1] / args.height
|
||||
x2 = boxesArr[j][2] / args.width
|
||||
y2 = boxesArr[j][3] / args.height
|
||||
left = int(x1 * frameWidth)
|
||||
top = int(y1 * frameHeight)
|
||||
width = int((x2 - x1) * frameWidth)
|
||||
height = int((y2 - y1) * frameHeight)
|
||||
classIds.append(int(classIdxArr[j]))
|
||||
confidences.append(score)
|
||||
boxes.append([left, top, width, height])
|
||||
elif len(outs) == 2 and outs[0].ndim == 4 and outs[0].shape[-1] == 4:
|
||||
boxesArr = outs[0].reshape(-1, 4)
|
||||
confsArr = outs[1].reshape(boxesArr.shape[0], -1)
|
||||
for j in range(boxesArr.shape[0]):
|
||||
classId = np.argmax(confsArr[j])
|
||||
confidence = float(confsArr[j][classId])
|
||||
if confidence > confThreshold:
|
||||
box = boxesArr[j]
|
||||
left = int(box[0] * frameWidth)
|
||||
top = int(box[1] * frameHeight)
|
||||
width = int((box[2] - box[0]) * frameWidth)
|
||||
height = int((box[3] - box[1]) * frameHeight)
|
||||
classIds.append(classId)
|
||||
confidences.append(confidence)
|
||||
boxes.append([left, top, width, height])
|
||||
else:
|
||||
print('Unsupported YOLO ONNX output format')
|
||||
exit()
|
||||
|
||||
elif args.postprocessing == 'yolov8' or args.postprocessing == 'yolov5':
|
||||
# Network produces output blob with a shape NxC where N is a number of
|
||||
# detected objects and C is a number of classes + 4 where the first 4
|
||||
# numbers are [center_x, center_y, width, height]
|
||||
box_scale_w = frameWidth / args.width
|
||||
box_scale_h = frameHeight / args.height
|
||||
|
||||
for out in outs:
|
||||
if args.postprocessing == 'yolov8':
|
||||
out = out[0].transpose(1, 0)
|
||||
else: # YOLOv5, no transposition needed
|
||||
out = out[0]
|
||||
|
||||
for detection in out:
|
||||
if args.postprocessing == 'yolov8':
|
||||
scores = detection[4:]
|
||||
obj_conf = 1
|
||||
else:
|
||||
scores = detection[5:]
|
||||
obj_conf = detection[4]
|
||||
|
||||
classId = np.argmax(scores)
|
||||
confidence = scores[classId]*obj_conf
|
||||
if confidence > confThreshold:
|
||||
center_x = int(detection[0] * box_scale_w)
|
||||
center_y = int(detection[1] * box_scale_h)
|
||||
width = int(detection[2] * box_scale_w)
|
||||
height = int(detection[3] * box_scale_h)
|
||||
left = int(center_x - width / 2)
|
||||
top = int(center_y - height / 2)
|
||||
classIds.append(classId)
|
||||
confidences.append(float(confidence))
|
||||
boxes.append([left, top, width, height])
|
||||
else:
|
||||
print('Unknown postprocessing method: ' + args.postprocessing)
|
||||
exit()
|
||||
|
||||
# NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample
|
||||
# or NMS is required if number of outputs > 1
|
||||
if len(outNames) > 1 or (args.postprocessing == 'yolov8' or args.postprocessing == 'yolov5') and args.backend != cv.dnn.DNN_BACKEND_OPENCV:
|
||||
indices = []
|
||||
classIds = np.array(classIds)
|
||||
boxes = np.array(boxes)
|
||||
confidences = np.array(confidences)
|
||||
unique_classes = set(classIds)
|
||||
for cl in unique_classes:
|
||||
class_indices = np.where(classIds == cl)[0]
|
||||
conf = confidences[class_indices]
|
||||
box = boxes[class_indices].tolist()
|
||||
nms_indices = cv.dnn.NMSBoxes(box, conf, confThreshold, nmsThreshold)
|
||||
indices.extend(class_indices[nms_indices])
|
||||
else:
|
||||
indices = np.arange(0, len(classIds))
|
||||
|
||||
return boxes, classIds, confidences, indices
|
||||
|
||||
def drawPred(classIds, confidences, boxes, indices, fontSize, fontThickness):
|
||||
for i in indices:
|
||||
box = boxes[i]
|
||||
left = box[0]
|
||||
top = box[1]
|
||||
right = box[0] + box[2]
|
||||
bottom = box[1] + box[3]
|
||||
bg_color = get_color(classIds[i])
|
||||
cv.rectangle(frame, (left, top), (right, bottom), bg_color, fontThickness)
|
||||
|
||||
label = '%.2f' % confidences[i]
|
||||
|
||||
# Print a label of class.
|
||||
if labels:
|
||||
assert(classIds[i] < len(labels))
|
||||
label = '%s: %s' % (labels[classIds[i]], label)
|
||||
|
||||
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
top = max(top, labelSize[1])
|
||||
cv.rectangle(frame, (int(left-fontThickness/2), top - labelSize[1]), (left + labelSize[0], top + baseLine), bg_color, cv.FILLED)
|
||||
cv.putText(frame, label, (left, top-fontThickness), cv.FONT_HERSHEY_SIMPLEX, fontSize, get_text_color(bg_color), fontThickness)
|
||||
|
||||
# Process inputs
|
||||
winName = 'Deep learning object detection in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_AUTOSIZE)
|
||||
|
||||
def callback(pos):
|
||||
global confThreshold
|
||||
confThreshold = pos / 100.0
|
||||
|
||||
cv.createTrackbar('Confidence threshold, %', winName, int(confThreshold * 100), 99, callback)
|
||||
|
||||
cap = cv.VideoCapture(cv.samples.findFileOrKeep(args.input) if args.input else 0)
|
||||
|
||||
class QueueFPS(queue.Queue):
|
||||
def __init__(self):
|
||||
queue.Queue.__init__(self)
|
||||
self.startTime = 0
|
||||
self.counter = 0
|
||||
|
||||
def put(self, v):
|
||||
queue.Queue.put(self, v)
|
||||
self.counter += 1
|
||||
if self.counter == 1:
|
||||
self.startTime = time.time()
|
||||
|
||||
def getFPS(self):
|
||||
return self.counter / (time.time() - self.startTime)
|
||||
|
||||
|
||||
process = True
|
||||
|
||||
#
|
||||
# Frames capturing thread
|
||||
#
|
||||
framesQueue = QueueFPS()
|
||||
def framesThreadBody():
|
||||
global framesQueue, process
|
||||
|
||||
while process:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
break
|
||||
framesQueue.put(frame)
|
||||
|
||||
|
||||
#
|
||||
# Frames processing thread
|
||||
#
|
||||
processedFramesQueue = queue.Queue()
|
||||
predictionsQueue = QueueFPS()
|
||||
def processingThreadBody():
|
||||
global processedFramesQueue, predictionsQueue, args, process, asyncN
|
||||
|
||||
futureOutputs = []
|
||||
while process:
|
||||
# Get a next frame
|
||||
frame = None
|
||||
try:
|
||||
frame = framesQueue.get_nowait()
|
||||
|
||||
if asyncN:
|
||||
if len(futureOutputs) == asyncN:
|
||||
frame = None # Skip the frame
|
||||
else:
|
||||
framesQueue.queue.clear() # Skip the rest of frames
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
|
||||
if not frame is None:
|
||||
frameHeight = frame.shape[0]
|
||||
frameWidth = frame.shape[1]
|
||||
|
||||
# Create a 4D blob from a frame.
|
||||
inpWidth = args.width if args.width else frameWidth
|
||||
inpHeight = args.height if args.height else frameHeight
|
||||
blob = cv.dnn.blobFromImage(frame, scalefactor=args.scale, mean=args.mean, size=(inpWidth, inpHeight), swapRB=args.rgb, ddepth=cv.CV_32F)
|
||||
processedFramesQueue.put(frame)
|
||||
|
||||
# Run a model
|
||||
net.setInput(blob)
|
||||
|
||||
if asyncN:
|
||||
futureOutputs.append(net.forwardAsync())
|
||||
else:
|
||||
outs = net.forward(outNames)
|
||||
net.printPerfProfile()
|
||||
predictionsQueue.put(copy.deepcopy(outs))
|
||||
|
||||
while futureOutputs and futureOutputs[0].wait_for(0):
|
||||
out = futureOutputs[0].get()
|
||||
predictionsQueue.put(copy.deepcopy([out]))
|
||||
|
||||
del futureOutputs[0]
|
||||
|
||||
if args.use_threads:
|
||||
framesThread = Thread(target=framesThreadBody)
|
||||
framesThread.start()
|
||||
|
||||
processingThread = Thread(target=processingThreadBody)
|
||||
processingThread.start()
|
||||
|
||||
#
|
||||
# Postprocessing and rendering loop
|
||||
#
|
||||
while cv.waitKey(1) < 0:
|
||||
try:
|
||||
# Request prediction first because they put after frames
|
||||
outs = predictionsQueue.get_nowait()
|
||||
frame = processedFramesQueue.get_nowait()
|
||||
imgWidth = max(frame.shape[:2])
|
||||
fontSize = (stdSize*imgWidth)/stdImgSize
|
||||
fontThickness = max(1,(stdWeight*imgWidth)//stdImgSize)
|
||||
|
||||
boxes, classIds, confidences, indices = postprocess(frame, outs)
|
||||
drawPred(classIds, confidences, boxes, indices, fontSize, fontThickness)
|
||||
fontSize = fontSize/2
|
||||
# Put efficiency information.
|
||||
if predictionsQueue.counter > 1:
|
||||
label = 'Camera: %.2f FPS' % (framesQueue.getFPS())
|
||||
cv.rectangle(frame, (0, 0), (int(260*fontSize), int(80*fontSize)), (255,255,255), cv.FILLED)
|
||||
cv.putText(frame, label, (0, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
|
||||
label = 'Network: %.2f FPS' % (predictionsQueue.getFPS())
|
||||
cv.putText(frame, label, (0, int(2*25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
|
||||
label = 'Skipped frames: %d' % (framesQueue.counter - predictionsQueue.counter)
|
||||
cv.putText(frame, label, (0, int(3*25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
|
||||
cv.imshow(winName, frame)
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
|
||||
process = False
|
||||
framesThread.join()
|
||||
processingThread.join()
|
||||
|
||||
else:
|
||||
# Non-threaded processing if --async is 0
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
cv.waitKey()
|
||||
break
|
||||
|
||||
frameHeight = frame.shape[0]
|
||||
frameWidth = frame.shape[1]
|
||||
|
||||
inpWidth = args.width if args.width else frameWidth
|
||||
inpHeight = args.height if args.height else frameHeight
|
||||
blob = cv.dnn.blobFromImage(frame, scalefactor=args.scale, mean=args.mean, size=(inpWidth, inpHeight), swapRB=args.rgb, ddepth=cv.CV_32F)
|
||||
|
||||
net.setInput(blob)
|
||||
outs = net.forward(outNames)
|
||||
net.printPerfProfile()
|
||||
|
||||
boxes, classIds, confidences, indices = postprocess(frame, outs)
|
||||
drawPred(classIds, confidences, boxes, indices, (stdSize*max(frame.shape[:2]))/stdImgSize, (stdWeight*max(frame.shape[:2]))//stdImgSize)
|
||||
|
||||
cv.imshow(winName, frame)
|
||||
@@ -0,0 +1,309 @@
|
||||
#include <iostream>
|
||||
#include <cmath>
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/video.hpp>
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
using namespace cv::dnn;
|
||||
|
||||
const string about = "Use this script for testing Object Tracking using OpenCV. \n\n"
|
||||
"Firstly, download required models using the download_models.py <alias>.\n"
|
||||
"Valid alias names are nanotrack, vit and dasiamrpn.\n\n"
|
||||
"To run:\n"
|
||||
"\t nanotrack: \n"
|
||||
"\t\t e.g: ./example_dnn_object_tracker nanotrack\n\n"
|
||||
"\t vit: \n"
|
||||
"\t\t e.g: ./example_dnn_object_tracker vit\n\n"
|
||||
"\t dasiamrpn: \n"
|
||||
"\t\t e.g: ./example_dnn_object_tracker dasiamrpn\n\n"
|
||||
|
||||
"To switch between models in runtime, make sure all the models are downloaded using download_models.py\n";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message }"
|
||||
"{ @alias | vit | 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 | | Full path to input video folder, the specific camera index. (empty for camera 0) }"
|
||||
"{ tracking_thrs | 0.3 | Tracking score threshold. If a bbox of score >= 0.3, it is considered as found }";
|
||||
|
||||
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;
|
||||
|
||||
static void loadParser(const string &modelName, const string &zooFile)
|
||||
{
|
||||
// Load appropriate preprocessing arguments based on model name
|
||||
if (modelName == "vit")
|
||||
{
|
||||
keys += genPreprocArguments(modelName, zooFile, "");
|
||||
}
|
||||
else if (modelName == "nanotrack")
|
||||
{
|
||||
keys += genPreprocArguments(modelName, zooFile, "nanotrack_head_");
|
||||
keys += genPreprocArguments(modelName, zooFile, "nanotrack_back_");
|
||||
}
|
||||
else if (modelName == "dasiamrpn")
|
||||
{
|
||||
keys += genPreprocArguments(modelName, zooFile, "dasiamrpn_");
|
||||
keys += genPreprocArguments(modelName, zooFile, "dasiamrpn_kernel_r1_");
|
||||
keys += genPreprocArguments(modelName, zooFile, "dasiamrpn_kernel_cls_");
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
|
||||
static void createTracker(const string &modelName, CommandLineParser &parser, Ptr<Tracker> &tracker) {
|
||||
int backend = getBackendID(parser.get<String>("backend"));
|
||||
int target = getTargetID(parser.get<String>("target"));
|
||||
if (modelName == "dasiamrpn") {
|
||||
const string net = parser.get<String>("dasiamrpn_model");
|
||||
const string sha1 = parser.get<String>("dasiamrpn_sha1");
|
||||
const string kernel_cls1 = parser.get<String>("dasiamrpn_kernel_cls_model");
|
||||
const string kernel_cls_sha1 = parser.get<String>("dasiamrpn_kernel_cls_sha1");
|
||||
const string kernel_r1 = parser.get<String>("dasiamrpn_kernel_r1_model");
|
||||
const string kernel_sha1 = parser.get<String>("dasiamrpn_kernel_r1_sha1");
|
||||
|
||||
TrackerDaSiamRPN::Params params;
|
||||
params.model = findModel(net, sha1);
|
||||
params.kernel_cls1 = findModel(kernel_cls1, kernel_cls_sha1);
|
||||
params.kernel_r1 = findModel(kernel_r1, kernel_sha1);
|
||||
params.backend = backend;
|
||||
params.target = target;
|
||||
tracker = TrackerDaSiamRPN::create(params);
|
||||
} else if (modelName == "nanotrack") {
|
||||
const string backbone = parser.get<String>("nanotrack_back_model");
|
||||
const string backSha1 = parser.get<String>("nanotrack_back_sha1");
|
||||
const string headneck = parser.get<String>("nanotrack_head_model");
|
||||
const string headSha1 = parser.get<String>("nanotrack_head_sha1");
|
||||
|
||||
TrackerNano::Params params;
|
||||
params.backbone = findModel(backbone, backSha1);
|
||||
params.neckhead = findModel(headneck, headSha1);
|
||||
params.backend = backend;
|
||||
params.target = target;
|
||||
tracker = TrackerNano::create(params);
|
||||
} else if (modelName == "vit") {
|
||||
const string net = parser.get<String>("model");
|
||||
const string sha1 = parser.get<String>("sha1");
|
||||
float tracking_score_threshold = parser.get<float>("tracking_thrs");
|
||||
|
||||
TrackerVit::Params params;
|
||||
params.net = findModel(net, sha1);
|
||||
params.backend = backend;
|
||||
params.target = target;
|
||||
params.tracking_score_threshold = tracking_score_threshold;
|
||||
tracker = TrackerVit::create(params);
|
||||
} else {
|
||||
cout<<"Pass the valid alias. Choices are {vit, nanotrack, dasiamrpn }."<<endl;
|
||||
exit(0);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
parser.about(about);
|
||||
if (!parser.has("@alias") || parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
string modelName = parser.get<String>("@alias");
|
||||
const string zooFile = findFile(parser.get<String>("zoo"));
|
||||
loadParser(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
|
||||
Ptr<Tracker> tracker;
|
||||
createTracker(modelName, parser, tracker);
|
||||
|
||||
const string windowName = "TRACKING";
|
||||
namedWindow(windowName, WINDOW_NORMAL);
|
||||
FontFace fontFace("sans");
|
||||
int stdSize = 20;
|
||||
int stdWeight = 400;
|
||||
int stdImgSize = 512;
|
||||
int imgWidth = -1;
|
||||
int fontSize = 50;
|
||||
int fontWeight = 500;
|
||||
double alpha = 0.4;
|
||||
Rect selectRect;
|
||||
string inputName = parser.get<String>("input");
|
||||
string instructionLabel = "Press space bar to pause video to draw bounding box.";
|
||||
Rect banner;
|
||||
// Open a video file or an image file or a camera stream.
|
||||
VideoCapture cap;
|
||||
|
||||
if (inputName.empty() || (isdigit(inputName[0]) && inputName.size() == 1))
|
||||
{
|
||||
int c = inputName.empty() ? 0 : inputName[0] - '0';
|
||||
cout << "Trying to open camera #" << c << " ..." << endl;
|
||||
if (!cap.open(c))
|
||||
{
|
||||
cout << "Capture from camera #" << c << " didn't work. Specify -i=<video> parameter to read from video file" << endl;
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
else if (inputName.size())
|
||||
{
|
||||
string filePath = findFile(inputName);
|
||||
if (!cap.open(filePath))
|
||||
{
|
||||
cout << "Could not open: " << inputName << endl;
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
Mat image;
|
||||
|
||||
for (;;)
|
||||
{
|
||||
cap >> image;
|
||||
if (image.empty())
|
||||
{
|
||||
cerr << "Can't capture frame. End of video stream?" << endl;
|
||||
return 0;
|
||||
}
|
||||
else if (imgWidth == -1){
|
||||
imgWidth = min(image.rows, image.cols);
|
||||
fontSize = (stdSize*imgWidth)/stdImgSize;
|
||||
fontWeight = (stdWeight*imgWidth)/stdImgSize;
|
||||
banner = getTextSize(Size(), instructionLabel, Point(), fontFace, fontSize, fontWeight);
|
||||
banner.height += 2 * fontSize; // padding
|
||||
banner.width += 10; // padding
|
||||
}
|
||||
Mat org_img = image.clone();
|
||||
rectangle(image, banner, Scalar::all(255), FILLED);
|
||||
addWeighted(image, alpha, org_img, 1 - alpha, 0, image);
|
||||
putText(image, instructionLabel, Point(10, fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
putText(image, "Press space bar after selecting.", Point(10, 2*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
imshow(windowName, image);
|
||||
int key = waitKey(30); //Simulating 30 FPS, if reduced frames move really fast
|
||||
if (key == ' ')
|
||||
{
|
||||
selectRect = selectROI(windowName, image);
|
||||
if (selectRect.width > 0 && selectRect.height > 0)
|
||||
{
|
||||
break;
|
||||
}
|
||||
else
|
||||
{
|
||||
cout << "No valid selection made. Please select again." << endl;
|
||||
}
|
||||
}
|
||||
else if (key == 27) // ESC key to exit
|
||||
{
|
||||
exit(0);
|
||||
}
|
||||
}
|
||||
|
||||
cout << "ROI=" << selectRect << endl;
|
||||
tracker->init(image, selectRect);
|
||||
instructionLabel = "Press space bar to select new target";
|
||||
banner = getTextSize(Size(), instructionLabel, Point(), fontFace, fontSize, fontWeight);
|
||||
banner.height += 4 * fontSize; // padding
|
||||
banner.width += 10; // padding
|
||||
|
||||
TickMeter tickMeter;
|
||||
|
||||
for (int count = 0; ; ++count)
|
||||
{
|
||||
cap >> image;
|
||||
if (image.empty())
|
||||
{
|
||||
cerr << "Can't capture frame " << count << ". End of video stream?" << endl;
|
||||
break;
|
||||
}
|
||||
Rect rect;
|
||||
|
||||
tickMeter.start();
|
||||
bool ok = tracker->update(image, rect);
|
||||
tickMeter.stop();
|
||||
|
||||
float score = tracker->getTrackingScore();
|
||||
|
||||
Mat render_image = image.clone();
|
||||
|
||||
int key = waitKey(30); //Simulating 30 FPS, if reduced frames move really fast
|
||||
int h = image.rows;
|
||||
int w = image.cols;
|
||||
rectangle(render_image, banner, Scalar::all(255), FILLED);
|
||||
rectangle(render_image, cv::Point(0, int(h - int(1.5*fontSize))), cv::Point(w, h), Scalar::all(255), FILLED);
|
||||
addWeighted(render_image, alpha, image, 1 - alpha, 0, render_image);
|
||||
putText(render_image, instructionLabel, Point(10, fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
putText(render_image, "For switching between trackers: press 'v' for ViT, 'n' for Nano, and 'd' for DaSiamRPN.", Point(10, h-10), Scalar(0,0,0), fontFace, int(0.8*fontSize), fontWeight);
|
||||
|
||||
if (ok){
|
||||
if (key == ' '){
|
||||
putText(render_image, "Select the new target", Point(10, 2*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
selectRect = selectROI(windowName, render_image);
|
||||
if (selectRect.width > 0 && selectRect.height > 0){
|
||||
tracker->init(image, selectRect);
|
||||
}
|
||||
else{
|
||||
cout<<"New target is not selected, switching to previous target"<<endl;
|
||||
}
|
||||
}
|
||||
else if (key == 'v'){
|
||||
modelName = "vit";
|
||||
loadParser(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
createTracker(modelName, parser, tracker);
|
||||
tracker->init(image, rect);
|
||||
}
|
||||
else if (key == 'n'){
|
||||
modelName = "nanotrack";
|
||||
loadParser(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
createTracker(modelName, parser, tracker);
|
||||
tracker->init(image, rect);
|
||||
}
|
||||
else if (key == 'd'){
|
||||
modelName = "dasiamrpn";
|
||||
loadParser(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
createTracker(modelName, parser, tracker);
|
||||
tracker->init(image, rect);
|
||||
}
|
||||
rectangle(render_image, rect, Scalar(0, 255, 0), 2);
|
||||
}
|
||||
|
||||
string timeLabel = format("FPS: %.2f", tickMeter.getFPS());
|
||||
string scoreLabel = format("Score: %f", score);
|
||||
string algoLabel = "Algorithm: " + modelName;
|
||||
putText(render_image, timeLabel, Point(10, 2*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
putText(render_image, scoreLabel, Point(10, 3*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
putText(render_image, algoLabel, Point(10, 4*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
|
||||
imshow(windowName, render_image);
|
||||
|
||||
tickMeter.reset();
|
||||
|
||||
if (key == 27 /*ESC*/)
|
||||
exit(0);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python
|
||||
import sys
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
from common import *
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Use this script for testing Object Tracking using OpenCV.
|
||||
Firstly, download required models using the download_models.py.
|
||||
To run:
|
||||
nanotrack:
|
||||
Download Model: python download_models.py nanotrack
|
||||
Example: python object_tracker.py nanotrack
|
||||
vit:
|
||||
Download Model: python download_models.py vit
|
||||
Example: python object_tracker.py vit
|
||||
or
|
||||
python object_tracker.py
|
||||
dasiamrpn:
|
||||
Download Model: python download_models.py dasiamrpn
|
||||
Example: python object_tracker.py dasiamrpn
|
||||
To switch between models in runtime, make sure all the models are downloaded using download_models.py'''
|
||||
)
|
||||
|
||||
def load_parser(model_name):
|
||||
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", type=str, help="Path to video source")
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
add_preproc_args(args.zoo, parser, 'object_tracker', alias=model_name)
|
||||
if model_name == "dasiamrpn":
|
||||
add_preproc_args(args.zoo, parser, 'object_tracker', prefix="dasiamrpn_", alias="dasiamrpn")
|
||||
add_preproc_args(args.zoo, parser, 'object_tracker', prefix="dasiamrpn_kernel_r1_", alias="dasiamrpn")
|
||||
add_preproc_args(args.zoo, parser, 'object_tracker', prefix="dasiamrpn_kernel_cls_", alias="dasiamrpn")
|
||||
elif model_name == "nanotrack":
|
||||
add_preproc_args(args.zoo, parser, 'object_tracker', prefix="nanotrack_back_", alias="nanotrack")
|
||||
add_preproc_args(args.zoo, parser, 'object_tracker', prefix="nanotrack_head_", alias="nanotrack")
|
||||
elif model_name != "vit":
|
||||
print("Pass the valid alias. Choices are { nanotrack, vit, dasiamrpn }")
|
||||
exit(0)
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='''
|
||||
Firstly, download required models using `python download_models.py {modelName}`
|
||||
Run using python object_tracker.py {modelName}.
|
||||
''',
|
||||
formatter_class=argparse.RawTextHelpFormatter)
|
||||
return parser.parse_args()
|
||||
|
||||
def createTracker(model_name, args):
|
||||
if model_name == 'dasiamrpn':
|
||||
print("Using Dasiamrpn Tracker.")
|
||||
params = cv.TrackerDaSiamRPN_Params()
|
||||
params.model = findModel(args.dasiamrpn_model, args.dasiamrpn_sha1)
|
||||
params.kernel_cls1 = findModel(args.dasiamrpn_kernel_cls_model, args.dasiamrpn_kernel_cls_sha1)
|
||||
params.kernel_r1 = findModel(args.dasiamrpn_kernel_r1_model, args.dasiamrpn_kernel_r1_sha1)
|
||||
tracker = cv.TrackerDaSiamRPN_create(params)
|
||||
elif model_name == 'nanotrack':
|
||||
print("Using Nano Tracker.")
|
||||
params = cv.TrackerNano_Params()
|
||||
params.backbone = findModel(args.nanotrack_back_model, args.nanotrack_back_sha1)
|
||||
params.neckhead = findModel(args.nanotrack_head_model, args.nanotrack_head_sha1)
|
||||
tracker = cv.TrackerNano_create(params)
|
||||
elif model_name == 'vit':
|
||||
print("Using Vit Tracker.")
|
||||
params = cv.TrackerVit_Params()
|
||||
params.net = findModel(args.model, args.sha1)
|
||||
tracker = cv.TrackerVit_create(params)
|
||||
else:
|
||||
help()
|
||||
exit(-1)
|
||||
return tracker
|
||||
|
||||
def main(model_name, args):
|
||||
tracker = createTracker(model_name, args)
|
||||
videoPath = args.input
|
||||
print('Using video: {}'.format(videoPath))
|
||||
cap = cv.VideoCapture(cv.samples.findFile(args.input) if args.input else 0)
|
||||
if not cap.isOpened():
|
||||
print("Can't open video stream: {}".format(videoPath))
|
||||
exit(-1)
|
||||
|
||||
stdSize = 0.6
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
imgWidth = -1 # Initialization
|
||||
fontSize = 1.5
|
||||
fontThickness = 1
|
||||
alpha = 0.5
|
||||
windowName = "TRACKING"
|
||||
cv.namedWindow(windowName, cv.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
ret, image = cap.read()
|
||||
if not ret:
|
||||
print("Video completed!!")
|
||||
return -1
|
||||
if imgWidth == -1:
|
||||
imgWidth = min(image.shape[:2])
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(fontThickness,(stdWeight*imgWidth)//stdImgSize)
|
||||
label = "Press space bar to pause video to draw bounding box."
|
||||
labelSize, _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
org_img = image.copy()
|
||||
cv.rectangle(image, (0, 0), (labelSize[0]+10, labelSize[1]+int(40*fontSize)), (255,255,255), cv.FILLED)
|
||||
cv.addWeighted(image, alpha, org_img, 1 - alpha, 0, image)
|
||||
cv.putText(image, label, (10, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.putText(image, "Press space bar after selecting.", (10, int(55*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.imshow(windowName, image)
|
||||
|
||||
key = cv.waitKey(30) & 0xFF
|
||||
if key == ord(' '):
|
||||
bbox = cv.selectROI(windowName, image)
|
||||
print('ROI: {}'.format(bbox))
|
||||
if bbox != (0, 0, 0, 0):
|
||||
break
|
||||
|
||||
if key == ord('q') or key == 27:
|
||||
return
|
||||
try:
|
||||
tracker.init(image, bbox)
|
||||
except Exception as e:
|
||||
print('Unable to initialize tracker with requested bounding box. Is there any object?')
|
||||
print(e)
|
||||
|
||||
tick_meter = cv.TickMeter()
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
if imgWidth == -1:
|
||||
imgWidth = min(frame.shape[:2])
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(fontThickness,(stdWeight*imgWidth)//stdImgSize)
|
||||
label="Press space bar to select new target"
|
||||
labelSize, _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
tick_meter.reset()
|
||||
tick_meter.start()
|
||||
ok, newbox = tracker.update(frame)
|
||||
tick_meter.stop()
|
||||
score = tracker.getTrackingScore()
|
||||
render_image = frame.copy()
|
||||
key = cv.waitKey(30) & 0xFF
|
||||
h, w = frame.shape[:2]
|
||||
cv.rectangle(render_image, (0, 0), (labelSize[0]+10, labelSize[1]+int(100*fontSize)), (255,255,255), cv.FILLED)
|
||||
cv.rectangle(render_image, (0, int(h-45*fontSize)), (w, h), (255,255,255), cv.FILLED)
|
||||
cv.addWeighted(render_image, alpha, frame, 1 - alpha, 0, render_image)
|
||||
cv.putText(render_image, label, (10, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.putText(render_image, "For switching between trackers: press 'v' for ViT, 'n' for Nanotrack, and 'd' for DaSiamRPN.", (10, h-10), cv.FONT_HERSHEY_SIMPLEX, 0.8*fontSize, (0, 0, 0), fontThickness)
|
||||
|
||||
if ok:
|
||||
if key == ord(' '):
|
||||
cv.putText(render_image, "Select the new target", (10, int(55*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
bbox = cv.selectROI(windowName, render_image)
|
||||
print('ROI:', bbox)
|
||||
if bbox != (0, 0, 0, 0):
|
||||
tracker.init(frame, bbox)
|
||||
elif key == ord('v'):
|
||||
model_name = "vit"
|
||||
args = load_parser(model_name)
|
||||
tracker = createTracker(model_name, args)
|
||||
tracker.init(frame, newbox)
|
||||
elif key == ord('n'):
|
||||
model_name = "nanotrack"
|
||||
args = load_parser(model_name)
|
||||
tracker = createTracker(model_name, args)
|
||||
tracker.init(frame, newbox)
|
||||
elif key == ord('d'):
|
||||
model_name = "dasiamrpn"
|
||||
args = load_parser(model_name)
|
||||
tracker = createTracker(model_name, args)
|
||||
tracker.init(frame, newbox)
|
||||
elif key == ord('q') or key == 27:
|
||||
return
|
||||
|
||||
cv.rectangle(render_image, newbox, (200, 0, 0), thickness=2)
|
||||
time_label = f"FPS: {tick_meter.getFPS():.2f}"
|
||||
score_label = f"Tracking score: {score:.2f}"
|
||||
algo_label = f"Algorithm: {model_name}"
|
||||
cv.putText(render_image, time_label, (10, int(55*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.putText(render_image, score_label, (10, int(85*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.putText(render_image, algo_label, (10, int(115*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
|
||||
cv.imshow(windowName, render_image)
|
||||
if key in [ord('q'), 27]:
|
||||
break
|
||||
|
||||
if __name__ == '__main__':
|
||||
help()
|
||||
if len(sys.argv) < 2 or sys.argv[1].startswith("--"):
|
||||
model_name = "vit"
|
||||
else:
|
||||
model_name = sys.argv[1]
|
||||
args = load_parser(model_name)
|
||||
|
||||
main(model_name, args)
|
||||
cv.destroyAllWindows()
|
||||
@@ -0,0 +1,157 @@
|
||||
//
|
||||
// this sample demonstrates the use of pretrained openpose networks with opencv's dnn module.
|
||||
//
|
||||
// it can be used for body pose detection, using either the COCO model(18 parts):
|
||||
// http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/coco/pose_iter_440000.caffemodel
|
||||
// https://raw.githubusercontent.com/opencv/opencv_extra/5.x/testdata/dnn/openpose_pose_coco.prototxt
|
||||
//
|
||||
// or the MPI model(16 parts):
|
||||
// http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/mpi/pose_iter_160000.caffemodel
|
||||
// https://raw.githubusercontent.com/opencv/opencv_extra/5.x/testdata/dnn/openpose_pose_mpi_faster_4_stages.prototxt
|
||||
//
|
||||
// (to simplify this sample, the body models are restricted to a single person.)
|
||||
//
|
||||
//
|
||||
// you can also try the hand pose model:
|
||||
// http://posefs1.perception.cs.cmu.edu/OpenPose/models/hand/pose_iter_102000.caffemodel
|
||||
// https://raw.githubusercontent.com/CMU-Perceptual-Computing-Lab/openpose/master/models/hand/pose_deploy.prototxt
|
||||
//
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
using namespace cv;
|
||||
using namespace cv::dnn;
|
||||
|
||||
#include <iostream>
|
||||
using namespace std;
|
||||
|
||||
|
||||
// connection table, in the format [model_id][pair_id][from/to]
|
||||
// please look at the nice explanation at the bottom of:
|
||||
// https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/output.md
|
||||
//
|
||||
const int POSE_PAIRS[3][20][2] = {
|
||||
{ // COCO body
|
||||
{1,2}, {1,5}, {2,3},
|
||||
{3,4}, {5,6}, {6,7},
|
||||
{1,8}, {8,9}, {9,10},
|
||||
{1,11}, {11,12}, {12,13},
|
||||
{1,0}, {0,14},
|
||||
{14,16}, {0,15}, {15,17}
|
||||
},
|
||||
{ // MPI body
|
||||
{0,1}, {1,2}, {2,3},
|
||||
{3,4}, {1,5}, {5,6},
|
||||
{6,7}, {1,14}, {14,8}, {8,9},
|
||||
{9,10}, {14,11}, {11,12}, {12,13}
|
||||
},
|
||||
{ // hand
|
||||
{0,1}, {1,2}, {2,3}, {3,4}, // thumb
|
||||
{0,5}, {5,6}, {6,7}, {7,8}, // pinkie
|
||||
{0,9}, {9,10}, {10,11}, {11,12}, // middle
|
||||
{0,13}, {13,14}, {14,15}, {15,16}, // ring
|
||||
{0,17}, {17,18}, {18,19}, {19,20} // small
|
||||
}};
|
||||
|
||||
int main(int argc, char **argv)
|
||||
{
|
||||
CommandLineParser parser(argc, argv,
|
||||
"{ h help | false | print this help message }"
|
||||
"{ p proto | | (required) model configuration, e.g. hand/pose.prototxt }"
|
||||
"{ m model | | (required) model weights, e.g. hand/pose_iter_102000.caffemodel }"
|
||||
"{ i image | | (required) path to image file (containing a single person, or hand) }"
|
||||
"{ d dataset | | specify what kind of model was trained. It could be (COCO, MPI, HAND) depends on dataset. }"
|
||||
"{ width | 368 | Preprocess input image by resizing to a specific width. }"
|
||||
"{ height | 368 | Preprocess input image by resizing to a specific height. }"
|
||||
"{ t threshold | 0.1 | threshold or confidence value for the heatmap }"
|
||||
"{ s scale | 0.003922 | scale for blob }"
|
||||
);
|
||||
|
||||
String modelTxt = samples::findFile(parser.get<string>("proto"));
|
||||
String modelBin = samples::findFile(parser.get<string>("model"));
|
||||
String imageFile = samples::findFile(parser.get<String>("image"));
|
||||
String dataset = parser.get<String>("dataset");
|
||||
int W_in = parser.get<int>("width");
|
||||
int H_in = parser.get<int>("height");
|
||||
float thresh = parser.get<float>("threshold");
|
||||
float scale = parser.get<float>("scale");
|
||||
|
||||
if (parser.get<bool>("help") || modelTxt.empty() || modelBin.empty() || imageFile.empty())
|
||||
{
|
||||
cout << "A sample app to demonstrate human or hand pose detection with a pretrained OpenPose dnn." << endl;
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
int midx, npairs, nparts;
|
||||
if (!dataset.compare("COCO")) { midx = 0; npairs = 17; nparts = 18; }
|
||||
else if (!dataset.compare("MPI")) { midx = 1; npairs = 14; nparts = 16; }
|
||||
else if (!dataset.compare("HAND")) { midx = 2; npairs = 20; nparts = 22; }
|
||||
else
|
||||
{
|
||||
std::cerr << "Can't interpret dataset parameter: " << dataset << std::endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
// read the network model
|
||||
Net net = readNet(modelBin, modelTxt);
|
||||
// and the image
|
||||
Mat img = imread(imageFile);
|
||||
if (img.empty())
|
||||
{
|
||||
std::cerr << "Can't read image from the file: " << imageFile << std::endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
// send it through the network
|
||||
Mat inputBlob = blobFromImage(img, scale, Size(W_in, H_in), Scalar(0, 0, 0), false, false);
|
||||
net.setInput(inputBlob);
|
||||
Mat result = net.forward();
|
||||
// the result is an array of "heatmaps", the probability of a body part being in location x,y
|
||||
|
||||
int H = result.size[2];
|
||||
int W = result.size[3];
|
||||
|
||||
// find the position of the body parts
|
||||
vector<Point> points(22);
|
||||
for (int n=0; n<nparts; n++)
|
||||
{
|
||||
// Slice heatmap of corresponding body's part.
|
||||
Mat heatMap(H, W, CV_32F, result.ptr(0,n));
|
||||
// 1 maximum per heatmap
|
||||
Point p(-1,-1),pm;
|
||||
double conf;
|
||||
minMaxLoc(heatMap, 0, &conf, 0, &pm);
|
||||
if (conf > thresh)
|
||||
p = pm;
|
||||
points[n] = p;
|
||||
}
|
||||
|
||||
// connect body parts and draw it !
|
||||
float SX = float(img.cols) / W;
|
||||
float SY = float(img.rows) / H;
|
||||
for (int n=0; n<npairs; n++)
|
||||
{
|
||||
// lookup 2 connected body/hand parts
|
||||
Point2f a = points[POSE_PAIRS[midx][n][0]];
|
||||
Point2f b = points[POSE_PAIRS[midx][n][1]];
|
||||
|
||||
// we did not find enough confidence before
|
||||
if (a.x<=0 || a.y<=0 || b.x<=0 || b.y<=0)
|
||||
continue;
|
||||
|
||||
// scale to image size
|
||||
a.x*=SX; a.y*=SY;
|
||||
b.x*=SX; b.y*=SY;
|
||||
|
||||
line(img, a, b, Scalar(0,200,0), 2);
|
||||
circle(img, a, 3, Scalar(0,0,200), -1);
|
||||
circle(img, b, 3, Scalar(0,0,200), -1);
|
||||
}
|
||||
|
||||
imshow("OpenPose", img);
|
||||
waitKey();
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
# To use Inference Engine backend, specify location of plugins:
|
||||
# source /opt/intel/computer_vision_sdk/bin/setupvars.sh
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description='This script is used to demonstrate OpenPose human pose estimation network '
|
||||
'from https://github.com/CMU-Perceptual-Computing-Lab/openpose project using OpenCV. '
|
||||
'The sample and model are simplified and could be used for a single person on the frame.')
|
||||
parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera')
|
||||
parser.add_argument('--proto', help='Path to .prototxt')
|
||||
parser.add_argument('--model', help='Path to .caffemodel')
|
||||
parser.add_argument('--dataset', help='Specify what kind of model was trained. '
|
||||
'It could be (COCO, MPI, HAND) depends on dataset.')
|
||||
parser.add_argument('--thr', default=0.1, type=float, help='Threshold value for pose parts heat map')
|
||||
parser.add_argument('--width', default=368, type=int, help='Resize input to specific width.')
|
||||
parser.add_argument('--height', default=368, type=int, help='Resize input to specific height.')
|
||||
parser.add_argument('--scale', default=0.003922, type=float, help='Scale for blob.')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dataset == 'COCO':
|
||||
BODY_PARTS = { "Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4,
|
||||
"LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9,
|
||||
"RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "REye": 14,
|
||||
"LEye": 15, "REar": 16, "LEar": 17, "Background": 18 }
|
||||
|
||||
POSE_PAIRS = [ ["Neck", "RShoulder"], ["Neck", "LShoulder"], ["RShoulder", "RElbow"],
|
||||
["RElbow", "RWrist"], ["LShoulder", "LElbow"], ["LElbow", "LWrist"],
|
||||
["Neck", "RHip"], ["RHip", "RKnee"], ["RKnee", "RAnkle"], ["Neck", "LHip"],
|
||||
["LHip", "LKnee"], ["LKnee", "LAnkle"], ["Neck", "Nose"], ["Nose", "REye"],
|
||||
["REye", "REar"], ["Nose", "LEye"], ["LEye", "LEar"] ]
|
||||
elif args.dataset == 'MPI':
|
||||
BODY_PARTS = { "Head": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4,
|
||||
"LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9,
|
||||
"RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "Chest": 14,
|
||||
"Background": 15 }
|
||||
|
||||
POSE_PAIRS = [ ["Head", "Neck"], ["Neck", "RShoulder"], ["RShoulder", "RElbow"],
|
||||
["RElbow", "RWrist"], ["Neck", "LShoulder"], ["LShoulder", "LElbow"],
|
||||
["LElbow", "LWrist"], ["Neck", "Chest"], ["Chest", "RHip"], ["RHip", "RKnee"],
|
||||
["RKnee", "RAnkle"], ["Chest", "LHip"], ["LHip", "LKnee"], ["LKnee", "LAnkle"] ]
|
||||
elif args.dataset == 'HAND':
|
||||
BODY_PARTS = { "Wrist": 0,
|
||||
"ThumbMetacarpal": 1, "ThumbProximal": 2, "ThumbMiddle": 3, "ThumbDistal": 4,
|
||||
"IndexFingerMetacarpal": 5, "IndexFingerProximal": 6, "IndexFingerMiddle": 7, "IndexFingerDistal": 8,
|
||||
"MiddleFingerMetacarpal": 9, "MiddleFingerProximal": 10, "MiddleFingerMiddle": 11, "MiddleFingerDistal": 12,
|
||||
"RingFingerMetacarpal": 13, "RingFingerProximal": 14, "RingFingerMiddle": 15, "RingFingerDistal": 16,
|
||||
"LittleFingerMetacarpal": 17, "LittleFingerProximal": 18, "LittleFingerMiddle": 19, "LittleFingerDistal": 20,
|
||||
}
|
||||
|
||||
POSE_PAIRS = [ ["Wrist", "ThumbMetacarpal"], ["ThumbMetacarpal", "ThumbProximal"],
|
||||
["ThumbProximal", "ThumbMiddle"], ["ThumbMiddle", "ThumbDistal"],
|
||||
["Wrist", "IndexFingerMetacarpal"], ["IndexFingerMetacarpal", "IndexFingerProximal"],
|
||||
["IndexFingerProximal", "IndexFingerMiddle"], ["IndexFingerMiddle", "IndexFingerDistal"],
|
||||
["Wrist", "MiddleFingerMetacarpal"], ["MiddleFingerMetacarpal", "MiddleFingerProximal"],
|
||||
["MiddleFingerProximal", "MiddleFingerMiddle"], ["MiddleFingerMiddle", "MiddleFingerDistal"],
|
||||
["Wrist", "RingFingerMetacarpal"], ["RingFingerMetacarpal", "RingFingerProximal"],
|
||||
["RingFingerProximal", "RingFingerMiddle"], ["RingFingerMiddle", "RingFingerDistal"],
|
||||
["Wrist", "LittleFingerMetacarpal"], ["LittleFingerMetacarpal", "LittleFingerProximal"],
|
||||
["LittleFingerProximal", "LittleFingerMiddle"], ["LittleFingerMiddle", "LittleFingerDistal"] ]
|
||||
else:
|
||||
raise(Exception("you need to specify either 'COCO', 'MPI', or 'Hand' in args.dataset"))
|
||||
|
||||
inWidth = args.width
|
||||
inHeight = args.height
|
||||
inScale = args.scale
|
||||
|
||||
net = cv.dnn.readNet(cv.samples.findFile(args.proto), cv.samples.findFile(args.model))
|
||||
|
||||
cap = cv.VideoCapture(args.input if args.input else 0)
|
||||
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
cv.waitKey()
|
||||
break
|
||||
|
||||
frameWidth = frame.shape[1]
|
||||
frameHeight = frame.shape[0]
|
||||
inp = cv.dnn.blobFromImage(frame, inScale, (inWidth, inHeight),
|
||||
(0, 0, 0), swapRB=False, crop=False)
|
||||
net.setInput(inp)
|
||||
t0 = cv.getTickCount()
|
||||
out = net.forward()
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
|
||||
assert(len(BODY_PARTS) <= out.shape[1])
|
||||
|
||||
points = []
|
||||
for i in range(len(BODY_PARTS)):
|
||||
# Slice heatmap of corresponding body's part.
|
||||
heatMap = out[0, i, :, :]
|
||||
|
||||
# Originally, we try to find all the local maximums. To simplify a sample
|
||||
# we just find a global one. However only a single pose at the same time
|
||||
# could be detected this way.
|
||||
_, conf, _, point = cv.minMaxLoc(heatMap)
|
||||
x = (frameWidth * point[0]) / out.shape[3]
|
||||
y = (frameHeight * point[1]) / out.shape[2]
|
||||
|
||||
# Add a point if it's confidence is higher than threshold.
|
||||
points.append((int(x), int(y)) if conf > args.thr else None)
|
||||
|
||||
for pair in POSE_PAIRS:
|
||||
partFrom = pair[0]
|
||||
partTo = pair[1]
|
||||
assert(partFrom in BODY_PARTS)
|
||||
assert(partTo in BODY_PARTS)
|
||||
|
||||
idFrom = BODY_PARTS[partFrom]
|
||||
idTo = BODY_PARTS[partTo]
|
||||
|
||||
if points[idFrom] and points[idTo]:
|
||||
cv.line(frame, points[idFrom], points[idTo], (0, 255, 0), 3)
|
||||
cv.ellipse(frame, points[idFrom], (3, 3), 0, 0, 360, (0, 0, 255), cv.FILLED)
|
||||
cv.ellipse(frame, points[idTo], (3, 3), 0, 0, 360, (0, 0, 255), cv.FILLED)
|
||||
|
||||
cv.putText(frame, '%.2f ms' % (t * 1000.0), (10, 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
|
||||
|
||||
cv.imshow('OpenPose using OpenCV', frame)
|
||||
@@ -0,0 +1,91 @@
|
||||
#!/usr/bin/env python
|
||||
'''
|
||||
This sample uses the RAFT model to calculate optical flow.
|
||||
|
||||
RAFT Original Paper: https://arxiv.org/pdf/2003.12039.pdf
|
||||
RAFT Repo: https://github.com/princeton-vl/RAFT
|
||||
|
||||
Download the .onnx model from here https://github.com/opencv/opencv_zoo/raw/281d232cd99cd920853106d853c440edd35eb442/models/optical_flow_estimation_raft/optical_flow_estimation_raft_2023aug.onnx.
|
||||
|
||||
Note: the legacy FlowNet v2 Caffe pipeline (--proto/.caffemodel) has been removed together
|
||||
with the Caffe importer. Please provide a single ONNX model.
|
||||
'''
|
||||
|
||||
import argparse
|
||||
import os.path
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
|
||||
|
||||
class OpticalFlow(object):
|
||||
def __init__(self, model, height, width, proto=""):
|
||||
if proto:
|
||||
raise cv.error("Caffe support has been removed. Please provide a single ONNX model path (e.g. RAFT).")
|
||||
self.net = cv.dnn.readNet(model)
|
||||
self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
|
||||
self.height = height
|
||||
self.width = width
|
||||
|
||||
def compute_flow(self, first_img, second_img):
|
||||
inp0 = cv.dnn.blobFromImage(first_img, size=(self.width, self.height))
|
||||
inp1 = cv.dnn.blobFromImage(second_img, size=(self.width, self.height))
|
||||
self.net.setInputsNames(["img0", "img1"])
|
||||
self.net.setInput(inp0, "img0")
|
||||
self.net.setInput(inp1, "img1")
|
||||
|
||||
flow = self.net.forward()
|
||||
output = self.motion_to_color(flow)
|
||||
return output
|
||||
|
||||
def motion_to_color(self, flow):
|
||||
arr = np.arange(0, 255, dtype=np.uint8)
|
||||
colormap = cv.applyColorMap(arr, cv.COLORMAP_HSV)
|
||||
colormap = colormap.squeeze(1)
|
||||
|
||||
flow = flow.squeeze(0)
|
||||
fx, fy = flow[0, ...], flow[1, ...]
|
||||
rad = np.sqrt(fx**2 + fy**2)
|
||||
maxrad = rad.max() if rad.max() != 0 else 1
|
||||
|
||||
ncols = arr.size
|
||||
rad = rad[..., np.newaxis] / maxrad
|
||||
a = np.arctan2(-fy / maxrad, -fx / maxrad) / np.pi
|
||||
fk = (a + 1) / 2.0 * (ncols - 1)
|
||||
k0 = fk.astype(np.int32)
|
||||
k1 = (k0 + 1) % ncols
|
||||
f = fk[..., np.newaxis] - k0[..., np.newaxis]
|
||||
|
||||
col0 = colormap[k0] / 255.0
|
||||
col1 = colormap[k1] / 255.0
|
||||
col = (1 - f) * col0 + f * col1
|
||||
col = np.where(rad <= 1, 1 - rad * (1 - col), col * 0.75)
|
||||
output = (255.0 * col).astype(np.uint8)
|
||||
return output
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='Use this script to calculate optical flow',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('-input', '-i', required=True, help='Path to input video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--height', default=320, type=int, help='Input height')
|
||||
parser.add_argument('--width', default=448, type=int, help='Input width')
|
||||
parser.add_argument('--model', '-m', required=True, help='Path to a single ONNX model (e.g. RAFT).')
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
if not os.path.isfile(args.model):
|
||||
raise OSError("Model does not exist")
|
||||
|
||||
winName = 'Calculation optical flow in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_NORMAL)
|
||||
cap = cv.VideoCapture(args.input if args.input else 0)
|
||||
hasFrame, first_frame = cap.read()
|
||||
|
||||
opt_flow = OpticalFlow(args.model, 360, 480)
|
||||
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, second_frame = cap.read()
|
||||
if not hasFrame:
|
||||
break
|
||||
flow = opt_flow.compute_flow(first_frame, second_frame)
|
||||
first_frame = second_frame
|
||||
cv.imshow(winName, flow)
|
||||
@@ -0,0 +1,380 @@
|
||||
/*
|
||||
This sample detects the query person in the given video file.
|
||||
|
||||
Authors of samples and Youtu ReID baseline:
|
||||
Xing Sun <winfredsun@tencent.com>
|
||||
Feng Zheng <zhengf@sustech.edu.cn>
|
||||
Xinyang Jiang <sevjiang@tencent.com>
|
||||
Fufu Yu <fufuyu@tencent.com>
|
||||
Enwei Zhang <miyozhang@tencent.com>
|
||||
|
||||
Copyright (C) 2020-2021, Tencent.
|
||||
Copyright (C) 2020-2021, SUSTech.
|
||||
Copyright (C) 2024, Bigvision LLC.
|
||||
|
||||
How to use:
|
||||
sample command to run:
|
||||
|
||||
./example_dnn_person_reid
|
||||
The system will ask you to mark the person to be tracked
|
||||
|
||||
You can download ReID model using:
|
||||
`python download_models.py reid`
|
||||
and yolo model using:
|
||||
`python download_models.py yolov8`
|
||||
|
||||
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 cv::dnn;
|
||||
using namespace std;
|
||||
|
||||
const string about = "Use this script for Person Re-identification using OpenCV. \n\n"
|
||||
"Firstly, download required models i.e. reid and yolov8 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_person_reid reid\n\n"
|
||||
"Re-Identification model path can also be specified using --model argument. Detection model can be set using --yolo_model argument.\n\n";
|
||||
|
||||
const string param_keys =
|
||||
"{help h | | show help message}"
|
||||
"{ @alias | reid | 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 }"
|
||||
"{query q | | Path to target image. Skip this argument to select target in the video frame.}"
|
||||
"{input i | | video file path}"
|
||||
"{yolo_model | | Path to yolov8n.onnx}";
|
||||
|
||||
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;
|
||||
|
||||
|
||||
struct MatComparator
|
||||
{
|
||||
bool operator()(const Mat &a, const Mat &b) const
|
||||
{
|
||||
return a.data < b.data; // This is a simple pointer comparison, not content!
|
||||
}
|
||||
};
|
||||
|
||||
map<Mat, Rect, MatComparator> imgDict;
|
||||
int height, width, yoloHeight, yoloWidth;
|
||||
float scale, yoloScale;
|
||||
bool swapRB, yoloSwapRB;
|
||||
Scalar mean_v, stnd;
|
||||
|
||||
|
||||
static void extractFeatures(vector<Mat> &imglist, Net &net, vector<Mat> &features)
|
||||
{
|
||||
for (size_t st = 0; st < imglist.size(); st++)
|
||||
{
|
||||
Mat blob;
|
||||
blobFromImage(imglist[st], blob, scale, Size(width, height), mean_v, swapRB, false, CV_32F);
|
||||
|
||||
// Check if standard deviation values are non-zero
|
||||
if (stnd[0] != 0.0 && stnd[1] != 0.0 && stnd[2] != 0.0)
|
||||
{
|
||||
// Divide blob by std for each channel
|
||||
divide(blob, stnd, blob);
|
||||
}
|
||||
net.setInput(blob);
|
||||
Mat out=net.forward();
|
||||
vector<int> s {out.size[0], out.size[1]};
|
||||
out = out.reshape(1, s);
|
||||
for (int i = 0; i < out.rows; i++)
|
||||
{
|
||||
Mat norm_features;
|
||||
normalize(out.row(i), norm_features, 1.0, 0.0, NORM_L2);
|
||||
features.push_back(norm_features);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
static int findMatching(const Mat &queryFeatures, const vector<Mat> &galleryFeatures)
|
||||
{
|
||||
if (queryFeatures.empty() || galleryFeatures.empty())
|
||||
return -1; // No valid index if either feature list is empty
|
||||
|
||||
int bestIndex = -1;
|
||||
float maxSimilarity = FLT_MIN;
|
||||
|
||||
for (int j = 0; j < (int)galleryFeatures.size(); j++)
|
||||
{
|
||||
float currentSimilarity = static_cast<float>(queryFeatures.dot(galleryFeatures[j]));
|
||||
if (currentSimilarity > maxSimilarity)
|
||||
{
|
||||
maxSimilarity = currentSimilarity;
|
||||
bestIndex = j;
|
||||
}
|
||||
}
|
||||
return bestIndex;
|
||||
}
|
||||
|
||||
static void yoloDetector(Mat &frame, Net &net, vector<Mat>& images)
|
||||
{
|
||||
int ht = frame.rows;
|
||||
int wt = frame.cols;
|
||||
|
||||
int length = max(ht, wt);
|
||||
|
||||
Mat image = Mat::zeros(Size(length, length), frame.type());
|
||||
|
||||
frame.copyTo(image(Rect(0, 0, wt, ht)));
|
||||
|
||||
// Calculate the scale
|
||||
double norm_scale = static_cast<double>(length) / yoloWidth;
|
||||
|
||||
Mat blob;
|
||||
blobFromImage(image, blob, yoloScale, Size(yoloWidth, yoloHeight), Scalar(), yoloSwapRB, false, CV_32F);
|
||||
net.setInput(blob);
|
||||
|
||||
vector<Mat> outputs;
|
||||
net.forward(outputs);
|
||||
Mat reshapedMatrix = outputs[0].reshape(0, 84); // Reshape to 2D (84 rows, 8400 columns)
|
||||
|
||||
Mat outputTransposed;
|
||||
transpose(reshapedMatrix, outputTransposed);
|
||||
|
||||
int rows = outputTransposed.rows;
|
||||
|
||||
vector<Rect2d> boxes;
|
||||
vector<float> scores;
|
||||
vector<int> class_ids;
|
||||
|
||||
for (int i = 0; i < rows; i++) {
|
||||
double minScore, maxScore;
|
||||
Point minClassLoc, maxClassLoc;
|
||||
minMaxLoc(outputTransposed.row(i).colRange(4, outputTransposed.cols), &minScore, &maxScore, &minClassLoc, &maxClassLoc);
|
||||
|
||||
if (maxScore >= 0.25 && maxClassLoc.x == 0) {
|
||||
double centerX = outputTransposed.at<float>(i, 0);
|
||||
double centerY = outputTransposed.at<float>(i, 1);
|
||||
double w = outputTransposed.at<float>(i, 2);
|
||||
double h = outputTransposed.at<float>(i, 3);
|
||||
|
||||
Rect2d box(
|
||||
centerX - 0.5 * w, // x
|
||||
centerY - 0.5 * h, // y
|
||||
w, // width
|
||||
h // height
|
||||
);
|
||||
boxes.push_back(box);
|
||||
scores.push_back(static_cast<float>(maxScore));
|
||||
class_ids.push_back(maxClassLoc.x); // x location gives the index
|
||||
}
|
||||
}
|
||||
|
||||
// Apply Non-Maximum Suppression
|
||||
vector<int> indexes;
|
||||
NMSBoxes(boxes, scores, 0.25f, 0.45f, indexes, 0.5f, 0);
|
||||
|
||||
images.resize(indexes.size());
|
||||
for (size_t i = 0; i < indexes.size(); i++) {
|
||||
int index = indexes[i];
|
||||
int x = static_cast<int>(round(boxes[index].x * norm_scale));
|
||||
int y = static_cast<int>(round(boxes[index].y * norm_scale));
|
||||
int w = static_cast<int>(round(boxes[index].width * norm_scale));
|
||||
int h = static_cast<int>(round(boxes[index].height * norm_scale));
|
||||
// Make sure the box is within the frame
|
||||
x = max(0, x);
|
||||
y = max(0, y);
|
||||
w = min(w, frame.cols - x);
|
||||
h = min(h, frame.rows - y);
|
||||
|
||||
// Crop the image
|
||||
Rect roi(x, y, w, h); // Define a region of interest
|
||||
images[i] = frame(roi); // Crop the region from the frame
|
||||
imgDict[images[i]] = roi;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
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);
|
||||
keys += genPreprocArguments(modelName, zooFile, "yolo_");
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
parser.about("Use this script to run ReID networks using OpenCV.");
|
||||
|
||||
const string sha1 = parser.get<String>("sha1");
|
||||
const string yoloSha1 = parser.get<String>("yolo_sha1");
|
||||
const string modelPath = findModel(parser.get<String>("model"), sha1);
|
||||
const string queryImagePath = parser.get<String>("query");
|
||||
string videoPath = parser.get<String>("input");
|
||||
const string yoloPath = findModel(parser.get<String>("yolo_model"), yoloSha1);
|
||||
const string backend = parser.get<String>("backend");
|
||||
const string target = parser.get<String>("target");
|
||||
height = parser.get<int>("height");
|
||||
width = parser.get<int>("width");
|
||||
yoloHeight = parser.get<int>("yolo_height");
|
||||
yoloWidth = parser.get<int>("yolo_width");
|
||||
scale = parser.get<float>("scale");
|
||||
yoloScale = parser.get<float>("yolo_scale");
|
||||
swapRB = parser.get<bool>("rgb");
|
||||
yoloSwapRB = parser.get<bool>("yolo_rgb");
|
||||
mean_v = parser.get<Scalar>("mean");
|
||||
stnd = parser.get<Scalar>("std");
|
||||
int stdSize = 20;
|
||||
int stdWeight = 400;
|
||||
int stdImgSize = 512;
|
||||
int imgWidth = -1; // Initialization
|
||||
int fontSize = 50;
|
||||
int fontWeight = 500;
|
||||
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if (backend != "default" || target != "cpu"){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
Net reidNet = readNetFromONNX(modelPath, engine);
|
||||
reidNet.setPreferableBackend(getBackendID(backend));
|
||||
reidNet.setPreferableTarget(getTargetID(target));
|
||||
|
||||
if(yoloPath.empty()){
|
||||
cout<<"[ERROR] Please pass path to yolov8.onnx model file using --yolo_model."<<endl;
|
||||
return -1;
|
||||
}
|
||||
Net net = readNetFromONNX(yoloPath, engine);
|
||||
|
||||
FontFace fontFace("sans");
|
||||
|
||||
VideoCapture cap;
|
||||
if (!videoPath.empty()){
|
||||
videoPath = findFile(videoPath);
|
||||
cap.open(videoPath);
|
||||
}
|
||||
else
|
||||
cap.open(0);
|
||||
|
||||
if (!cap.isOpened()) {
|
||||
cerr << "Error: Video could not be opened." << endl;
|
||||
return -1;
|
||||
}
|
||||
vector<Mat> queryImages;
|
||||
Mat queryImg;
|
||||
if (!queryImagePath.empty()) {
|
||||
queryImg = imread(queryImagePath);
|
||||
if (queryImg.empty()) {
|
||||
cerr << "Error: Query image could not be loaded." << endl;
|
||||
return -1;
|
||||
}
|
||||
queryImages.push_back(queryImg);
|
||||
} else {
|
||||
Mat image;
|
||||
for(;;) {
|
||||
cap.read(image);
|
||||
if (image.empty()) {
|
||||
cerr << "Error reading the video" << endl;
|
||||
return -1;
|
||||
}
|
||||
if (imgWidth == -1){
|
||||
imgWidth = min(image.rows, image.cols);
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize);
|
||||
fontWeight = min(fontWeight, (stdWeight*imgWidth)/stdImgSize);
|
||||
}
|
||||
|
||||
const string label = "Press space bar to pause video to draw bounding box.";
|
||||
Rect r = getTextSize(Size(), label, Point(), fontFace, fontSize, fontWeight);
|
||||
r.height += 2 * fontSize; // padding
|
||||
r.width += 10; // padding
|
||||
rectangle(image, r, Scalar::all(255), FILLED);
|
||||
putText(image, label, Point(10, fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
putText(image, "Press space bar after selecting.", Point(10, 2*fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
imshow("TRACKING", image);
|
||||
int key = waitKey(200);
|
||||
if(key == ' '){
|
||||
Rect rect = selectROI("TRACKING", image);
|
||||
|
||||
if (rect.width > 0 && rect.height > 0) {
|
||||
queryImg = image(rect).clone();
|
||||
queryImages.push_back(queryImg);
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (key == 'q' || key == 27) {
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Mat frame;
|
||||
vector<Mat> queryFeatures;
|
||||
extractFeatures(queryImages, reidNet, queryFeatures);
|
||||
|
||||
vector<Mat> detectedImages;
|
||||
vector<Mat> galleryFeatures;
|
||||
for(;;) {
|
||||
if (!cap.read(frame) || frame.empty()) {
|
||||
break;
|
||||
}
|
||||
if (imgWidth == -1){
|
||||
imgWidth = min(frame.rows, frame.cols);
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize);
|
||||
fontWeight = min(fontWeight, (stdWeight*imgWidth)/stdImgSize);
|
||||
}
|
||||
detectedImages.clear();
|
||||
galleryFeatures.clear();
|
||||
|
||||
yoloDetector(frame, net, detectedImages);
|
||||
extractFeatures(detectedImages, reidNet, galleryFeatures);
|
||||
|
||||
int match_idx = findMatching(queryFeatures[0], galleryFeatures);
|
||||
if (match_idx != -1 && static_cast<int>(detectedImages.size()) > match_idx) {
|
||||
Mat matchImg = detectedImages[match_idx];
|
||||
Rect bbox = imgDict[matchImg];
|
||||
rectangle(frame, bbox, Scalar(0, 0, 255), 2);
|
||||
putText(frame, "Target", Point(bbox.x, bbox.y - 10), Scalar(0,0,255), fontFace, fontSize, fontWeight);
|
||||
}
|
||||
const string label = "Tracking";
|
||||
Rect r = getTextSize(Size(), label, Point(), fontFace, fontSize, fontWeight);
|
||||
r.height += fontSize; // padding
|
||||
r.width += 10; // padding
|
||||
rectangle(frame, r, Scalar::all(255), FILLED);
|
||||
putText(frame, label, Point(10, fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
imshow("TRACKING", frame);
|
||||
int key = waitKey(30);
|
||||
if (key == 'q' || key == 27) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
cap.release();
|
||||
destroyAllWindows();
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
#!/usr/bin/env python
|
||||
'''
|
||||
This sample detects the query person in the given video file.
|
||||
|
||||
Authors of samples and Youtu ReID baseline:
|
||||
Xing Sun <winfredsun@tencent.com>
|
||||
Feng Zheng <zhengf@sustech.edu.cn>
|
||||
Xinyang Jiang <sevjiang@tencent.com>
|
||||
Fufu Yu <fufuyu@tencent.com>
|
||||
Enwei Zhang <miyozhang@tencent.com>
|
||||
|
||||
Copyright (C) 2020-2021, Tencent.
|
||||
Copyright (C) 2020-2021, SUSTech.
|
||||
Copyright (C) 2024, Bigvision LLC.
|
||||
|
||||
How to use:
|
||||
sample command to run:
|
||||
`python person_reid.py`
|
||||
|
||||
You can download ReID model using
|
||||
`python download_models.py reid`
|
||||
and yolo model using:
|
||||
`python download_models.py yolov8`
|
||||
|
||||
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 Person Re-identification using OpenCV.
|
||||
|
||||
Firstly, download required models i.e. reid and yolov8 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 person_reid.py reid
|
||||
|
||||
Re-identification model path can also be specified using --model argument and detection model can be specified using --yolo_model argument.
|
||||
'''
|
||||
)
|
||||
|
||||
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('--query', '-q', help='Path to target image. Skip this argument to select target in the video frame.')
|
||||
parser.add_argument('--input', '-i', default=0, help='Path to video 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, 'person_reid', prefix="", alias="reid")
|
||||
add_preproc_args(args.zoo, parser, 'person_reid', prefix="yolo_", alias="reid")
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Person Re-identification using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args()
|
||||
|
||||
img_dict = {} # Dictionary to store bounding boxes for corresponding cropped image
|
||||
|
||||
def yolo_detector(frame, net):
|
||||
global img_dict
|
||||
height, width, _ = frame.shape
|
||||
|
||||
length = max((height, width))
|
||||
image = np.zeros((length, length, 3), np.uint8)
|
||||
image[0:height, 0:width] = frame
|
||||
|
||||
scale = length/args.yolo_width
|
||||
# Create blob from the frame with correct scale factor and size for the model
|
||||
|
||||
blob = cv.dnn.blobFromImage(image, scalefactor=args.yolo_scale, size=(args.yolo_width, args.yolo_height), swapRB=args.yolo_rgb)
|
||||
net.setInput(blob)
|
||||
outputs = net.forward()
|
||||
|
||||
outputs = np.array([cv.transpose(outputs[0])])
|
||||
rows = outputs.shape[1]
|
||||
|
||||
boxes = []
|
||||
scores = []
|
||||
class_ids = []
|
||||
|
||||
for i in range(rows):
|
||||
classes_scores = outputs[0][i][4:]
|
||||
(_, maxScore, _, (x, maxClassIndex)) = cv.minMaxLoc(classes_scores)
|
||||
if maxScore >= 0.25:
|
||||
box = [
|
||||
outputs[0][i][0] - (0.5 * outputs[0][i][2]),
|
||||
outputs[0][i][1] - (0.5 * outputs[0][i][3]),
|
||||
outputs[0][i][2],
|
||||
outputs[0][i][3],
|
||||
]
|
||||
boxes.append(box)
|
||||
scores.append(maxScore)
|
||||
class_ids.append(maxClassIndex)
|
||||
|
||||
# Apply Non-Maximum Suppression
|
||||
indexes = cv.dnn.NMSBoxes(boxes, scores, 0.25, 0.45, 0.5)
|
||||
|
||||
images = []
|
||||
for i in indexes:
|
||||
x, y, w, h = boxes[i]
|
||||
x = round(x*scale)
|
||||
y = round(y*scale)
|
||||
w = round(w*scale)
|
||||
h = round(h*scale)
|
||||
|
||||
x, y = max(0, x), max(0, y)
|
||||
w, h = min(w, frame.shape[1] - x), min(h, frame.shape[0] - y)
|
||||
crop_img = frame[y:y+h, x:x+w]
|
||||
images.append(crop_img)
|
||||
img_dict[crop_img.tobytes()] = (x, y, w, h)
|
||||
return images
|
||||
|
||||
def extract_feature(images, net):
|
||||
"""
|
||||
Extract features from images
|
||||
:param images: the input images
|
||||
:param net: the model network
|
||||
"""
|
||||
feat_list = []
|
||||
# net = reid_net.copy()
|
||||
for img in images:
|
||||
blob = cv.dnn.blobFromImage(img, scalefactor=args.scale, size=(args.width, args.height), mean=args.mean, swapRB=args.rgb, crop=False, ddepth=cv.CV_32F)
|
||||
|
||||
for j in range(blob.shape[1]):
|
||||
blob[:, j, :, :] /= args.std[j]
|
||||
|
||||
net.setInput(blob)
|
||||
feat = net.forward()
|
||||
feat = np.reshape(feat, (feat.shape[0], feat.shape[1]))
|
||||
feat_list.append(feat)
|
||||
|
||||
feats = np.concatenate(feat_list, axis = 0)
|
||||
return feats
|
||||
|
||||
def find_matching(query_feat, gallery_feat):
|
||||
"""
|
||||
Return the index of the gallery image most similar to the query image
|
||||
:param query_feat: array of feature vectors of query images
|
||||
:param gallery_feat: array of feature vectors of gallery images
|
||||
"""
|
||||
cv.normalize(query_feat, query_feat, 1.0, 0.0, cv.NORM_L2)
|
||||
cv.normalize(gallery_feat, gallery_feat, 1.0, 0.0, cv.NORM_L2)
|
||||
|
||||
sim = query_feat.dot(gallery_feat.T)
|
||||
index = np.argmax(sim, axis=1)[0]
|
||||
return index
|
||||
|
||||
def main():
|
||||
if hasattr(args, 'help'):
|
||||
help()
|
||||
exit(1)
|
||||
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
|
||||
if args.yolo_model is None:
|
||||
print("[ERROR] Please pass path to yolov8.onnx model file using --yolo_model.")
|
||||
exit(1)
|
||||
else:
|
||||
args.yolo_model = findModel(args.yolo_model, args.yolo_sha1)
|
||||
|
||||
engine = cv.dnn.ENGINE_AUTO
|
||||
|
||||
if args.backend != "default" or args.target != "cpu":
|
||||
engine = cv.dnn.ENGINE_CLASSIC
|
||||
yolo_net = cv.dnn.readNetFromONNX(args.yolo_model, engine)
|
||||
reid_net = cv.dnn.readNetFromONNX(args.model, engine)
|
||||
reid_net.setPreferableBackend(get_backend_id(args.backend))
|
||||
reid_net.setPreferableTarget(get_target_id(args.target))
|
||||
cap = cv.VideoCapture(cv.samples.findFile(args.input) if args.input else 0)
|
||||
query_images = []
|
||||
|
||||
stdSize = 0.6
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
imgWidth = -1 # Initialization
|
||||
fontSize = 1.5
|
||||
fontThickness = 1
|
||||
|
||||
if args.query:
|
||||
query_images = [cv.imread(findFile(args.query))]
|
||||
else:
|
||||
while True:
|
||||
ret, image = cap.read()
|
||||
if not ret:
|
||||
print("Error reading the video")
|
||||
return -1
|
||||
if imgWidth == -1:
|
||||
imgWidth = min(image.shape[:2])
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(fontThickness,(stdWeight*imgWidth)//stdImgSize)
|
||||
|
||||
label = "Press space bar to pause video to draw bounding box."
|
||||
labelSize, _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
cv.rectangle(image, (0, 0), (labelSize[0]+10, labelSize[1]+int(30*fontSize)), (255,255,255), cv.FILLED)
|
||||
cv.putText(image, label, (10, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.putText(image, "Press space bar after selecting.", (10, int(50*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.imshow('TRACKING', image)
|
||||
|
||||
key = cv.waitKey(100) & 0xFF
|
||||
if key == ord(' '):
|
||||
rect = cv.selectROI("TRACKING", image)
|
||||
if rect:
|
||||
x, y, w, h = rect
|
||||
query_image = image[y:y + h, x:x + w]
|
||||
query_images = [query_image]
|
||||
break
|
||||
|
||||
if key == ord('q') or key == 27:
|
||||
return
|
||||
|
||||
query_feat = extract_feature(query_images, reid_net)
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
if imgWidth == -1:
|
||||
imgWidth = min(frame.shape[:2])
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(fontThickness,(stdWeight*imgWidth)//stdImgSize)
|
||||
|
||||
images = yolo_detector(frame, yolo_net)
|
||||
gallery_feat = extract_feature(images, reid_net)
|
||||
|
||||
match_idx = find_matching(query_feat, gallery_feat)
|
||||
|
||||
match_img = images[match_idx]
|
||||
x, y, w, h = img_dict[match_img.tobytes()]
|
||||
cv.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
|
||||
cv.putText(frame, "Target", (x, y - 10), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 255), fontThickness)
|
||||
|
||||
label="Tracking"
|
||||
labelSize, _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
cv.rectangle(frame, (0, 0), (labelSize[0]+10, labelSize[1]+10), (255,255,255), cv.FILLED)
|
||||
cv.putText(frame, label, (10, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
cv.imshow("TRACKING", frame)
|
||||
if cv.waitKey(1) & 0xFF in [ord('q'), 27]:
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv.destroyAllWindows()
|
||||
return
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args_parser()
|
||||
main()
|
||||
@@ -0,0 +1,134 @@
|
||||
'''
|
||||
This is a sample script to run Qwen2.5 inference in OpenCV using ONNX model.
|
||||
The script loads the Qwen2.5 model and runs inference on a given prompt using
|
||||
the ChatML format (<|im_start|> / <|im_end|> special tokens).
|
||||
|
||||
Model: https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct
|
||||
|
||||
Exporting Qwen2.5 model to ONNX:
|
||||
|
||||
1. Install the required dependencies:
|
||||
|
||||
pip install optimum[exporters] optimum-onnx[onnxruntime] torch transformers
|
||||
|
||||
2. Export the model to ONNX:
|
||||
|
||||
Without KV-cache:
|
||||
|
||||
optimum-cli export onnx --model Qwen/Qwen2.5-0.5B-Instruct --task causal-lm qwen2.5_instruct_onnx/
|
||||
|
||||
With KV-cache (recommended, faster autoregressive inference):
|
||||
|
||||
optimum-cli export onnx --model Qwen/Qwen2.5-0.5B-Instruct --task causal-lm-with-past qwen2.5_instruct_onnx_with_past/
|
||||
|
||||
|
||||
Run the script:
|
||||
1. Install the required dependencies:
|
||||
|
||||
pip install numpy
|
||||
|
||||
2. Run the script:
|
||||
|
||||
Without KV-cache (causal-lm export):
|
||||
|
||||
python qwen_inference.py --model=<path-to-onnx-model> \
|
||||
--tokenizer_path=<path-to-qwen2.5-config.json> \
|
||||
--prompt="What is OpenCV?"
|
||||
|
||||
With KV-cache (causal-lm-with-past export):
|
||||
|
||||
python qwen_inference.py --model=<path-to-onnx-model> \
|
||||
--tokenizer_path=<path-to-qwen2.5-config.json> \
|
||||
--prompt="What is OpenCV?" \
|
||||
--use_kv_cache
|
||||
'''
|
||||
|
||||
import numpy as np
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Use this script to run Qwen2.5 inference in OpenCV',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--model', type=str, required=True, help='Path to Qwen2.5 ONNX model file.')
|
||||
parser.add_argument('--tokenizer_path', type=str, required=True, help='Path to Qwen2.5 tokenizer config.json.')
|
||||
parser.add_argument('--prompt', type=str, default='What is OpenCV?', help='User prompt.')
|
||||
parser.add_argument('--max_new_tokens', type=int, default=64, help='Maximum number of new tokens to generate.')
|
||||
parser.add_argument('--use_kv_cache', action='store_true', default=False, help='Enable KV-cache for faster inference (requires causal-lm-with-past export).')
|
||||
parser.add_argument('--seed', type=int, default=0, help='Random seed.')
|
||||
return parser.parse_args()
|
||||
|
||||
def build_chatml_prompt(user_prompt):
|
||||
'''Wrap user prompt in Qwen2.5 ChatML format.'''
|
||||
return '<|im_start|>user\n' + user_prompt + '<|im_end|>\n<|im_start|>assistant\n'
|
||||
|
||||
def qwen_inference(net, prompt, max_new_tokens, tokenizer, use_kv_cache=True):
|
||||
|
||||
print("Inferencing Qwen2.5 model...")
|
||||
|
||||
tokens = list(tokenizer.encode(prompt))
|
||||
input_ids = np.array(tokens, dtype=np.int64).reshape(1, -1)
|
||||
|
||||
# Qwen2.5 special token IDs
|
||||
im_end_id = 151645 # <|im_end|>
|
||||
eos_id = 151643 # <|endoftext|>
|
||||
stop_ids = (im_end_id, eos_id)
|
||||
|
||||
generated = []
|
||||
|
||||
if use_kv_cache:
|
||||
net.enableKVCache()
|
||||
prompt_len = input_ids.shape[1]
|
||||
|
||||
# Prefill: process full prompt once to populate KV-cache
|
||||
net.setInput(input_ids, 'input_ids')
|
||||
net.setInput(np.ones((1, prompt_len), dtype=np.int64), 'attention_mask')
|
||||
net.setInput(np.arange(prompt_len, dtype=np.int64).reshape(1, -1), 'position_ids')
|
||||
logits = net.forward()
|
||||
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
|
||||
generated = [new_id]
|
||||
|
||||
# Generate: feed one new token per step; OpenCV routes present.* -> past_key_values.*
|
||||
for _ in range(max_new_tokens - 1):
|
||||
if new_id in stop_ids:
|
||||
break
|
||||
cur_len = prompt_len + len(generated)
|
||||
net.setInput(np.array([[new_id]], dtype=np.int64), 'input_ids')
|
||||
net.setInput(np.ones((1, cur_len), dtype=np.int64), 'attention_mask')
|
||||
net.setInput(np.array([[cur_len - 1]], dtype=np.int64), 'position_ids')
|
||||
logits = net.forward()
|
||||
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
|
||||
generated.append(new_id)
|
||||
else:
|
||||
# Without KV-cache: feed full growing sequence each step
|
||||
for _ in range(max_new_tokens):
|
||||
seq_len = input_ids.shape[1]
|
||||
net.setInput(input_ids, 'input_ids')
|
||||
net.setInput(np.ones((1, seq_len), dtype=np.int64), 'attention_mask')
|
||||
net.setInput(np.arange(seq_len, dtype=np.int64).reshape(1, -1), 'position_ids')
|
||||
logits = net.forward()
|
||||
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
|
||||
if new_id in stop_ids:
|
||||
break
|
||||
generated.append(new_id)
|
||||
input_ids = np.concatenate([input_ids, [[new_id]]], axis=1)
|
||||
|
||||
return np.array([tokens + generated], dtype=np.int64)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parse_args()
|
||||
np.random.seed(args.seed)
|
||||
|
||||
print("Preparing Qwen2.5 model...")
|
||||
tokenizer = cv.dnn.Tokenizer.load(args.tokenizer_path)
|
||||
|
||||
net = cv.dnn.readNetFromONNX(args.model, cv.dnn.ENGINE_NEW)
|
||||
|
||||
chatml_prompt = build_chatml_prompt(args.prompt)
|
||||
print(f"Prompt:\n{chatml_prompt}")
|
||||
|
||||
prompt_len = len(tokenizer.encode(chatml_prompt))
|
||||
tokens = qwen_inference(net, chatml_prompt, args.max_new_tokens, tokenizer, args.use_kv_cache)
|
||||
response = tokenizer.decode(tokens[0][prompt_len:].tolist())
|
||||
print(f"Response:\n{response}")
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 47 KiB |
@@ -0,0 +1,314 @@
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
#include <iostream>
|
||||
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/core/utils/logger.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
using namespace dnn;
|
||||
|
||||
const string about =
|
||||
"Use this script to run semantic segmentation deep learning networks using OpenCV.\n\n"
|
||||
"Firstly, download required models 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.\n"
|
||||
"To run:\n"
|
||||
"\t ./example_dnn_classification modelName(e.g. u2netp) --input=$OPENCV_SAMPLES_DATA_PATH/butterfly.jpg (or ignore this argument to use device camera)\n"
|
||||
"Model path can also be specified using --model argument.";
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ @alias | | 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 }"
|
||||
"{ device | 0 | camera device number. }"
|
||||
"{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }"
|
||||
"{ colors | | Optional path to a text file with colors for an every class. "
|
||||
"Every color is represented with three values from 0 to 255 in BGR channels order. }";
|
||||
|
||||
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;
|
||||
vector<string> labels;
|
||||
vector<Vec3b> colors;
|
||||
|
||||
|
||||
static void colorizeSegmentation(const Mat &score, Mat &segm)
|
||||
{
|
||||
const int rows = score.size[2];
|
||||
const int cols = score.size[3];
|
||||
const int chns = score.size[1];
|
||||
|
||||
if (colors.empty())
|
||||
{
|
||||
// Generate colors.
|
||||
colors.push_back(Vec3b());
|
||||
for (int i = 1; i < chns; ++i)
|
||||
{
|
||||
Vec3b color;
|
||||
for (int j = 0; j < 3; ++j)
|
||||
color[j] = (colors[i - 1][j] + rand() % 256) / 2;
|
||||
colors.push_back(color);
|
||||
}
|
||||
}
|
||||
else if (chns != (int)colors.size())
|
||||
{
|
||||
CV_Error(Error::StsError, format("Number of output labels does not match "
|
||||
"number of colors (%d != %zu)",
|
||||
chns, colors.size()));
|
||||
}
|
||||
|
||||
Mat maxCl = Mat::zeros(rows, cols, CV_8UC1);
|
||||
Mat maxVal(rows, cols, CV_32FC1, score.data);
|
||||
for (int ch = 1; ch < chns; ch++)
|
||||
{
|
||||
for (int row = 0; row < rows; row++)
|
||||
{
|
||||
const float *ptrScore = score.ptr<float>(0, ch, row);
|
||||
uint8_t *ptrMaxCl = maxCl.ptr<uint8_t>(row);
|
||||
float *ptrMaxVal = maxVal.ptr<float>(row);
|
||||
for (int col = 0; col < cols; col++)
|
||||
{
|
||||
if (ptrScore[col] > ptrMaxVal[col])
|
||||
{
|
||||
ptrMaxVal[col] = ptrScore[col];
|
||||
ptrMaxCl[col] = (uchar)ch;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
segm.create(rows, cols, CV_8UC3);
|
||||
for (int row = 0; row < rows; row++)
|
||||
{
|
||||
const uchar *ptrMaxCl = maxCl.ptr<uchar>(row);
|
||||
Vec3b *ptrSegm = segm.ptr<Vec3b>(row);
|
||||
for (int col = 0; col < cols; col++)
|
||||
{
|
||||
ptrSegm[col] = colors[ptrMaxCl[col]];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void showLegend(FontFace fontFace)
|
||||
{
|
||||
static const int kBlockHeight = 30;
|
||||
static Mat legend;
|
||||
if (legend.empty())
|
||||
{
|
||||
const int numClasses = (int)labels.size();
|
||||
if ((int)colors.size() != numClasses)
|
||||
{
|
||||
CV_Error(Error::StsError, format("Number of output labels does not match "
|
||||
"number of labels (%zu != %zu)",
|
||||
colors.size(), labels.size()));
|
||||
}
|
||||
legend.create(kBlockHeight * numClasses, 200, CV_8UC3);
|
||||
for (int i = 0; i < numClasses; i++)
|
||||
{
|
||||
Mat block = legend.rowRange(i * kBlockHeight, (i + 1) * kBlockHeight);
|
||||
block.setTo(colors[i]);
|
||||
Rect r = getTextSize(Size(), labels[i], Point(), fontFace, 15, 400);
|
||||
r.height += 15; // padding
|
||||
r.width += 10; // padding
|
||||
rectangle(block, r, Scalar::all(255), FILLED);
|
||||
putText(block, labels[i], Point(10, kBlockHeight/2), Scalar(0,0,0), fontFace, 15, 400);
|
||||
}
|
||||
namedWindow("Legend", WINDOW_AUTOSIZE);
|
||||
imshow("Legend", legend);
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char **argv)
|
||||
{
|
||||
utils::logging::setLogLevel(utils::logging::LOG_LEVEL_INFO);
|
||||
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
|
||||
const string modelName = parser.get<String>("@alias");
|
||||
const string zooFile = findFile(parser.get<String>("zoo"));
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
parser.about(about);
|
||||
if (!parser.has("@alias") || parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
string sha1 = parser.get<String>("sha1");
|
||||
float scale = parser.get<float>("scale");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
int inpWidth = parser.get<int>("width");
|
||||
int inpHeight = parser.get<int>("height");
|
||||
String model = findModel(parser.get<String>("model"), sha1);
|
||||
const string backend = parser.get<String>("backend");
|
||||
const string target = parser.get<String>("target");
|
||||
int stdSize = 20;
|
||||
int stdWeight = 400;
|
||||
int stdImgSize = 512;
|
||||
int imgWidth = -1; // Initialization
|
||||
int fontSize = 50;
|
||||
int fontWeight = 500;
|
||||
FontFace fontFace("sans");
|
||||
|
||||
// Open file with labels names.
|
||||
if (parser.has("labels"))
|
||||
{
|
||||
string file = findFile(parser.get<String>("labels"));
|
||||
ifstream ifs(file.c_str());
|
||||
if (!ifs.is_open())
|
||||
CV_Error(Error::StsError, "File " + file + " not found");
|
||||
string line;
|
||||
while (getline(ifs, line))
|
||||
{
|
||||
labels.push_back(line);
|
||||
}
|
||||
}
|
||||
// Open file with colors.
|
||||
if (parser.has("colors"))
|
||||
{
|
||||
string file = findFile(parser.get<String>("colors"));
|
||||
ifstream ifs(file.c_str());
|
||||
if (!ifs.is_open())
|
||||
CV_Error(Error::StsError, "File " + file + " not found");
|
||||
string line;
|
||||
while (getline(ifs, line))
|
||||
{
|
||||
istringstream colorStr(line.c_str());
|
||||
|
||||
Vec3b color;
|
||||
for (int i = 0; i < 3 && !colorStr.eof(); ++i)
|
||||
colorStr >> color[i];
|
||||
colors.push_back(color);
|
||||
}
|
||||
}
|
||||
|
||||
if (!parser.check())
|
||||
{
|
||||
parser.printErrors();
|
||||
return 1;
|
||||
}
|
||||
|
||||
CV_Assert(!model.empty());
|
||||
//! [Read and initialize network]
|
||||
EngineType engine = ENGINE_AUTO;
|
||||
if (backend != "default" || target != "cpu"){
|
||||
engine = ENGINE_CLASSIC;
|
||||
}
|
||||
Net net = readNetFromONNX(model, engine);
|
||||
net.setPreferableBackend(getBackendID(backend));
|
||||
net.setPreferableTarget(getTargetID(target));
|
||||
net.setProfilingMode(DNN_PROFILE_SUMMARY);
|
||||
//! [Read and initialize network]
|
||||
// Create a window
|
||||
static const string kWinName = "Deep learning semantic segmentation in OpenCV";
|
||||
namedWindow(kWinName, WINDOW_AUTOSIZE);
|
||||
|
||||
//! [Open a video file or an image file or a camera stream]
|
||||
VideoCapture cap;
|
||||
if (parser.has("input"))
|
||||
cap.open(findFile(parser.get<String>("input")));
|
||||
else
|
||||
cap.open(parser.get<int>("device"));
|
||||
|
||||
if (!cap.isOpened()) {
|
||||
cerr << "Error: Video could not be opened." << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
//! [Open a video file or an image file or a camera stream]
|
||||
// Process frames.
|
||||
Mat frame, blob;
|
||||
while (waitKey(1) < 0)
|
||||
{
|
||||
cap >> frame;
|
||||
if (frame.empty())
|
||||
{
|
||||
waitKey();
|
||||
break;
|
||||
}
|
||||
if (imgWidth == -1){
|
||||
imgWidth = max(frame.rows, frame.cols);
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize);
|
||||
fontWeight = min(fontWeight, (stdWeight*imgWidth)/stdImgSize);
|
||||
}
|
||||
imshow("Original Image", frame);
|
||||
//! [Create a 4D blob from a frame]
|
||||
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, false);
|
||||
//! [Set input blob]
|
||||
net.setInput(blob);
|
||||
//! [Set input blob]
|
||||
int64 t0 = getTickCount();
|
||||
|
||||
if (modelName == "u2netp")
|
||||
{
|
||||
vector<Mat> output;
|
||||
net.forward(output, net.getUnconnectedOutLayersNames());
|
||||
net.printPerfProfile();
|
||||
|
||||
Mat pred = output[0].reshape(1, output[0].size[2]);
|
||||
pred.convertTo(pred, CV_8U, 255.0);
|
||||
Mat mask;
|
||||
resize(pred, mask, Size(frame.cols, frame.rows), 0, 0, INTER_AREA);
|
||||
|
||||
// Create overlays for foreground and background
|
||||
Mat foreground_overlay;
|
||||
|
||||
// Set foreground (object) to red
|
||||
Mat all_zeros = Mat::zeros(frame.size(), CV_8UC1);
|
||||
vector<Mat> channels = {all_zeros, all_zeros, mask};
|
||||
merge(channels, foreground_overlay);
|
||||
|
||||
// Blend the overlays with the original frame
|
||||
addWeighted(frame, 0.25, foreground_overlay, 0.75, 0, frame);
|
||||
}
|
||||
else
|
||||
{
|
||||
//! [Make forward pass]
|
||||
Mat score = net.forward();
|
||||
net.printPerfProfile();
|
||||
//! [Make forward pass]
|
||||
Mat segm;
|
||||
colorizeSegmentation(score, segm);
|
||||
resize(segm, segm, frame.size(), 0, 0, INTER_NEAREST);
|
||||
addWeighted(frame, 0.1, segm, 0.9, 0.0, frame);
|
||||
}
|
||||
|
||||
// Put efficiency information.
|
||||
double t = (getTickCount() - t0) * 1000.0 / getTickFrequency();
|
||||
string label = format("Inference time: %.2f ms", t);
|
||||
Rect r = getTextSize(Size(), label, Point(), fontFace, fontSize, fontWeight);
|
||||
r.height += fontSize; // padding
|
||||
r.width += 10; // padding
|
||||
rectangle(frame, r, Scalar::all(255), FILLED);
|
||||
putText(frame, label, Point(10, fontSize), Scalar(0,0,0), fontFace, fontSize, fontWeight);
|
||||
|
||||
imshow(kWinName, frame);
|
||||
if (!labels.empty())
|
||||
showLegend(fontFace);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from common import *
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Firstly, download required models 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.\n"\n
|
||||
|
||||
To run:
|
||||
python segmentation.py model_name(e.g. u2netp) --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)
|
||||
|
||||
Model path can also be specified using --model argument
|
||||
'''
|
||||
)
|
||||
|
||||
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', help='Path to input image or video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--colors', help='Optional path to a text file with colors for an every class. '
|
||||
'An every color is represented with three values from 0 to 255 in BGR channels order.')
|
||||
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, 'segmentation')
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Use this script to run semantic segmentation deep learning networks using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args(func_args)
|
||||
|
||||
def showLegend(labels, colors, legend):
|
||||
if not labels is None and legend is None:
|
||||
blockHeight = 30
|
||||
assert(len(labels) == len(colors))
|
||||
|
||||
legend = np.zeros((blockHeight * len(colors), 200, 3), np.uint8)
|
||||
for i in range(len(labels)):
|
||||
block = legend[i * blockHeight:(i + 1) * blockHeight]
|
||||
block[:,:] = colors[i]
|
||||
cv.putText(block, labels[i], (0, blockHeight//2), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
|
||||
|
||||
cv.namedWindow('Legend', cv.WINDOW_AUTOSIZE)
|
||||
cv.imshow('Legend', legend)
|
||||
labels = None
|
||||
|
||||
def main(func_args=None):
|
||||
args = get_args_parser(func_args)
|
||||
if args.alias is None or hasattr(args, 'help'):
|
||||
help()
|
||||
exit(1)
|
||||
|
||||
cv.utils.logging.setLogLevel(cv.utils.logging.LOG_LEVEL_INFO)
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
if args.labels is not None:
|
||||
args.labels = findFile(args.labels)
|
||||
|
||||
np.random.seed(324)
|
||||
|
||||
stdSize = 0.8
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
imgWidth = -1 # Initialization
|
||||
fontSize = 1.5
|
||||
fontThickness = 1
|
||||
|
||||
# Load names of labels
|
||||
labels = None
|
||||
if args.labels:
|
||||
with open(args.labels, 'rt') as f:
|
||||
labels = f.read().rstrip('\n').split('\n')
|
||||
|
||||
# Load colors
|
||||
colors = None
|
||||
if args.colors:
|
||||
with open(args.colors, 'rt') as f:
|
||||
colors = [np.array(color.split(' '), np.uint8) for color in f.read().rstrip('\n').split('\n')]
|
||||
|
||||
# Load a network
|
||||
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))
|
||||
if hasattr(cv.dnn, 'DNN_PROFILE_SUMMARY'):
|
||||
net.setProfilingMode(cv.dnn.DNN_PROFILE_SUMMARY)
|
||||
|
||||
winName = 'Deep learning semantic segmentation in OpenCV'
|
||||
cv.namedWindow(winName, cv.WINDOW_AUTOSIZE)
|
||||
|
||||
cap = cv.VideoCapture(cv.samples.findFile(args.input) if args.input else 0)
|
||||
if not cap.isOpened():
|
||||
print("Failed to open the input video")
|
||||
exit(-1)
|
||||
|
||||
legend = None
|
||||
while cv.waitKey(1) < 0:
|
||||
hasFrame, frame = cap.read()
|
||||
if not hasFrame:
|
||||
cv.waitKey()
|
||||
break
|
||||
if imgWidth == -1:
|
||||
imgWidth = max(frame.shape[:2])
|
||||
fontSize = min(fontSize, (stdSize*imgWidth)/stdImgSize)
|
||||
fontThickness = max(fontThickness,(stdWeight*imgWidth)//stdImgSize)
|
||||
|
||||
cv.imshow("Original Image", frame)
|
||||
frameHeight = frame.shape[0]
|
||||
frameWidth = frame.shape[1]
|
||||
# Create a 4D blob from a frame.
|
||||
inpWidth = args.width if args.width else frameWidth
|
||||
inpHeight = args.height if args.height else frameHeight
|
||||
|
||||
blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=False)
|
||||
net.setInput(blob)
|
||||
|
||||
t0 = cv.getTickCount()
|
||||
if args.alias == 'u2netp':
|
||||
output = net.forward(net.getUnconnectedOutLayersNames())
|
||||
net.printPerfProfile()
|
||||
pred = output[0][0, 0, :, :]
|
||||
mask = (pred * 255).astype(np.uint8)
|
||||
mask = cv.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv.INTER_AREA)
|
||||
# Create overlays for foreground and background
|
||||
foreground_overlay = np.zeros_like(frame, dtype=np.uint8)
|
||||
# Set foreground (object) to red and background to blue
|
||||
foreground_overlay[:, :, 2] = mask # Red foreground
|
||||
# Blend the overlays with the original frame
|
||||
frame = cv.addWeighted(frame, 0.25, foreground_overlay, 0.75, 0)
|
||||
else:
|
||||
score = net.forward()
|
||||
net.printPerfProfile()
|
||||
|
||||
numClasses = score.shape[1]
|
||||
height = score.shape[2]
|
||||
width = score.shape[3]
|
||||
# Draw segmentation
|
||||
if not colors:
|
||||
# Generate colors
|
||||
colors = [np.array([0, 0, 0], np.uint8)]
|
||||
for i in range(1, numClasses):
|
||||
colors.append((colors[i - 1] + np.random.randint(0, 256, [3], np.uint8)) / 2)
|
||||
classIds = np.argmax(score[0], axis=0)
|
||||
segm = np.stack([colors[idx] for idx in classIds.flatten()])
|
||||
segm = segm.reshape(height, width, 3)
|
||||
|
||||
segm = cv.resize(segm, (frameWidth, frameHeight), interpolation=cv.INTER_NEAREST)
|
||||
frame = (0.1 * frame + 0.9 * segm).astype(np.uint8)
|
||||
|
||||
showLegend(labels, colors, legend)
|
||||
|
||||
label = 'Inference time: %.2f ms' % ((cv.getTickCount() - t0) * 1000.0 / cv.getTickFrequency())
|
||||
labelSize, _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, fontSize, fontThickness)
|
||||
cv.rectangle(frame, (0, 0), (labelSize[0]+10, labelSize[1]), (255,255,255), cv.FILLED)
|
||||
cv.putText(frame, label, (10, int(25*fontSize)), cv.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
|
||||
cv.imshow(winName, frame)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,62 @@
|
||||
# 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) 2017, Intel Corporation, all rights reserved.
|
||||
# Third party copyrights are property of their respective owners.
|
||||
import tensorflow as tf
|
||||
import struct
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
parser = argparse.ArgumentParser(description='Convert weights of a frozen TensorFlow graph to fp16.')
|
||||
parser.add_argument('--input', required=True, help='Path to frozen graph.')
|
||||
parser.add_argument('--output', required=True, help='Path to output graph.')
|
||||
parser.add_argument('--ops', default=['Conv2D', 'MatMul'], nargs='+',
|
||||
help='List of ops which weights are converted.')
|
||||
args = parser.parse_args()
|
||||
|
||||
DT_FLOAT = 1
|
||||
DT_HALF = 19
|
||||
|
||||
# For the frozen graphs, an every node that uses weights connected to Const nodes
|
||||
# through an Identity node. Usually they're called in the same way with '/read' suffix.
|
||||
# We'll replace all of them to Cast nodes.
|
||||
|
||||
# Load the model
|
||||
with tf.gfile.FastGFile(args.input) as f:
|
||||
graph_def = tf.GraphDef()
|
||||
graph_def.ParseFromString(f.read())
|
||||
|
||||
# Set of all inputs from desired nodes.
|
||||
inputs = []
|
||||
for node in graph_def.node:
|
||||
if node.op in args.ops:
|
||||
inputs += node.input
|
||||
|
||||
weightsNodes = []
|
||||
for node in graph_def.node:
|
||||
# From the whole inputs we need to keep only an Identity nodes.
|
||||
if node.name in inputs and node.op == 'Identity' and node.attr['T'].type == DT_FLOAT:
|
||||
weightsNodes.append(node.input[0])
|
||||
|
||||
# Replace Identity to Cast.
|
||||
node.op = 'Cast'
|
||||
node.attr['DstT'].type = DT_FLOAT
|
||||
node.attr['SrcT'].type = DT_HALF
|
||||
del node.attr['T']
|
||||
del node.attr['_class']
|
||||
|
||||
# Convert weights to halfs.
|
||||
for node in graph_def.node:
|
||||
if node.name in weightsNodes:
|
||||
node.attr['dtype'].type = DT_HALF
|
||||
node.attr['value'].tensor.dtype = DT_HALF
|
||||
|
||||
floats = node.attr['value'].tensor.tensor_content
|
||||
|
||||
floats = struct.unpack('f' * (len(floats) / 4), floats)
|
||||
halfs = np.array(floats).astype(np.float16).view(np.uint16)
|
||||
node.attr['value'].tensor.tensor_content = struct.pack('H' * len(halfs), *halfs)
|
||||
|
||||
tf.train.write_graph(graph_def, "", args.output, as_text=False)
|
||||
@@ -0,0 +1,403 @@
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
import numpy as np
|
||||
import os
|
||||
|
||||
"""
|
||||
Link to original paper : https://arxiv.org/abs/1812.11703
|
||||
Link to original repo : https://github.com/STVIR/pysot
|
||||
|
||||
You can download the pre-trained weights of the Tracker Model from https://drive.google.com/file/d/11bwgPFVkps9AH2NOD1zBDdpF_tQghAB-/view?usp=sharing
|
||||
You can download the target net (target branch of SiamRPN++) from https://drive.google.com/file/d/1dw_Ne3UMcCnFsaD6xkZepwE4GEpqq7U_/view?usp=sharing
|
||||
You can download the search net (search branch of SiamRPN++) from https://drive.google.com/file/d/1Lt4oE43ZSucJvze3Y-Z87CVDreO-Afwl/view?usp=sharing
|
||||
You can download the head model (RPN Head) from https://drive.google.com/file/d/1zT1yu12mtj3JQEkkfKFJWiZ71fJ-dQTi/view?usp=sharing
|
||||
"""
|
||||
|
||||
class ModelBuilder():
|
||||
""" This class generates the SiamRPN++ Tracker Model by using Imported ONNX Nets
|
||||
"""
|
||||
def __init__(self, target_net, search_net, rpn_head):
|
||||
super(ModelBuilder, self).__init__()
|
||||
# Build the target branch
|
||||
self.target_net = target_net
|
||||
# Build the search branch
|
||||
self.search_net = search_net
|
||||
# Build RPN_Head
|
||||
self.rpn_head = rpn_head
|
||||
|
||||
def template(self, z):
|
||||
""" Takes the template of size (1, 1, 127, 127) as an input to generate kernel
|
||||
"""
|
||||
self.target_net.setInput(z)
|
||||
outNames = self.target_net.getUnconnectedOutLayersNames()
|
||||
self.zfs_1, self.zfs_2, self.zfs_3 = self.target_net.forward(outNames)
|
||||
|
||||
def track(self, x):
|
||||
""" Takes the search of size (1, 1, 255, 255) as an input to generate classification score and bounding box regression
|
||||
"""
|
||||
self.search_net.setInput(x)
|
||||
outNames = self.search_net.getUnconnectedOutLayersNames()
|
||||
xfs_1, xfs_2, xfs_3 = self.search_net.forward(outNames)
|
||||
self.rpn_head.setInput(np.stack([self.zfs_1, self.zfs_2, self.zfs_3]), 'input_1')
|
||||
self.rpn_head.setInput(np.stack([xfs_1, xfs_2, xfs_3]), 'input_2')
|
||||
outNames = self.rpn_head.getUnconnectedOutLayersNames()
|
||||
cls, loc = self.rpn_head.forward(outNames)
|
||||
return {'cls': cls, 'loc': loc}
|
||||
|
||||
class Anchors:
|
||||
""" This class generate anchors.
|
||||
"""
|
||||
def __init__(self, stride, ratios, scales, image_center=0, size=0):
|
||||
self.stride = stride
|
||||
self.ratios = ratios
|
||||
self.scales = scales
|
||||
self.image_center = image_center
|
||||
self.size = size
|
||||
self.anchor_num = len(self.scales) * len(self.ratios)
|
||||
self.anchors = self.generate_anchors()
|
||||
|
||||
def generate_anchors(self):
|
||||
"""
|
||||
generate anchors based on predefined configuration
|
||||
"""
|
||||
anchors = np.zeros((self.anchor_num, 4), dtype=np.float32)
|
||||
size = self.stride**2
|
||||
count = 0
|
||||
for r in self.ratios:
|
||||
ws = int(np.sqrt(size * 1. / r))
|
||||
hs = int(ws * r)
|
||||
|
||||
for s in self.scales:
|
||||
w = ws * s
|
||||
h = hs * s
|
||||
anchors[count][:] = [-w * 0.5, -h * 0.5, w * 0.5, h * 0.5][:]
|
||||
count += 1
|
||||
return anchors
|
||||
|
||||
class SiamRPNTracker:
|
||||
def __init__(self, model):
|
||||
super(SiamRPNTracker, self).__init__()
|
||||
self.anchor_stride = 8
|
||||
self.anchor_ratios = [0.33, 0.5, 1, 2, 3]
|
||||
self.anchor_scales = [8]
|
||||
self.track_base_size = 8
|
||||
self.track_context_amount = 0.5
|
||||
self.track_exemplar_size = 127
|
||||
self.track_instance_size = 255
|
||||
self.track_lr = 0.4
|
||||
self.track_penalty_k = 0.04
|
||||
self.track_window_influence = 0.44
|
||||
self.score_size = (self.track_instance_size - self.track_exemplar_size) // \
|
||||
self.anchor_stride + 1 + self.track_base_size
|
||||
self.anchor_num = len(self.anchor_ratios) * len(self.anchor_scales)
|
||||
hanning = np.hanning(self.score_size)
|
||||
window = np.outer(hanning, hanning)
|
||||
self.window = np.tile(window.flatten(), self.anchor_num)
|
||||
self.anchors = self.generate_anchor(self.score_size)
|
||||
self.model = model
|
||||
|
||||
def get_subwindow(self, im, pos, model_sz, original_sz, avg_chans):
|
||||
"""
|
||||
Args:
|
||||
im: bgr based input image frame
|
||||
pos: position of the center of the frame
|
||||
model_sz: exemplar / target image size
|
||||
s_z: original / search image size
|
||||
avg_chans: channel average
|
||||
Return:
|
||||
im_patch: sub_windows for the given image input
|
||||
"""
|
||||
if isinstance(pos, float):
|
||||
pos = [pos, pos]
|
||||
sz = original_sz
|
||||
im_h, im_w, im_d = im.shape
|
||||
c = (original_sz + 1) / 2
|
||||
cx, cy = pos
|
||||
context_xmin = np.floor(cx - c + 0.5)
|
||||
context_xmax = context_xmin + sz - 1
|
||||
context_ymin = np.floor(cy - c + 0.5)
|
||||
context_ymax = context_ymin + sz - 1
|
||||
left_pad = int(max(0., -context_xmin))
|
||||
top_pad = int(max(0., -context_ymin))
|
||||
right_pad = int(max(0., context_xmax - im_w + 1))
|
||||
bottom_pad = int(max(0., context_ymax - im_h + 1))
|
||||
context_xmin += left_pad
|
||||
context_xmax += left_pad
|
||||
context_ymin += top_pad
|
||||
context_ymax += top_pad
|
||||
|
||||
if any([top_pad, bottom_pad, left_pad, right_pad]):
|
||||
size = (im_h + top_pad + bottom_pad, im_w + left_pad + right_pad, im_d)
|
||||
te_im = np.zeros(size, np.uint8)
|
||||
te_im[top_pad:top_pad + im_h, left_pad:left_pad + im_w, :] = im
|
||||
if top_pad:
|
||||
te_im[0:top_pad, left_pad:left_pad + im_w, :] = avg_chans
|
||||
if bottom_pad:
|
||||
te_im[im_h + top_pad:, left_pad:left_pad + im_w, :] = avg_chans
|
||||
if left_pad:
|
||||
te_im[:, 0:left_pad, :] = avg_chans
|
||||
if right_pad:
|
||||
te_im[:, im_w + left_pad:, :] = avg_chans
|
||||
im_patch = te_im[int(context_ymin):int(context_ymax + 1),
|
||||
int(context_xmin):int(context_xmax + 1), :]
|
||||
else:
|
||||
im_patch = im[int(context_ymin):int(context_ymax + 1),
|
||||
int(context_xmin):int(context_xmax + 1), :]
|
||||
|
||||
if not np.array_equal(model_sz, original_sz):
|
||||
im_patch = cv.resize(im_patch, (model_sz, model_sz))
|
||||
im_patch = im_patch.transpose(2, 0, 1)
|
||||
im_patch = im_patch[np.newaxis, :, :, :]
|
||||
im_patch = im_patch.astype(np.float32)
|
||||
return im_patch
|
||||
|
||||
def generate_anchor(self, score_size):
|
||||
"""
|
||||
Args:
|
||||
im: bgr based input image frame
|
||||
pos: position of the center of the frame
|
||||
model_sz: exemplar / target image size
|
||||
s_z: original / search image size
|
||||
avg_chans: channel average
|
||||
Return:
|
||||
anchor: anchors for pre-determined values of stride, ratio, and scale
|
||||
"""
|
||||
anchors = Anchors(self.anchor_stride, self.anchor_ratios, self.anchor_scales)
|
||||
anchor = anchors.anchors
|
||||
x1, y1, x2, y2 = anchor[:, 0], anchor[:, 1], anchor[:, 2], anchor[:, 3]
|
||||
anchor = np.stack([(x1 + x2) * 0.5, (y1 + y2) * 0.5, x2 - x1, y2 - y1], 1)
|
||||
total_stride = anchors.stride
|
||||
anchor_num = anchors.anchor_num
|
||||
anchor = np.tile(anchor, score_size * score_size).reshape((-1, 4))
|
||||
ori = - (score_size // 2) * total_stride
|
||||
xx, yy = np.meshgrid([ori + total_stride * dx for dx in range(score_size)],
|
||||
[ori + total_stride * dy for dy in range(score_size)])
|
||||
xx, yy = np.tile(xx.flatten(), (anchor_num, 1)).flatten(), \
|
||||
np.tile(yy.flatten(), (anchor_num, 1)).flatten()
|
||||
anchor[:, 0], anchor[:, 1] = xx.astype(np.float32), yy.astype(np.float32)
|
||||
return anchor
|
||||
|
||||
def _convert_bbox(self, delta, anchor):
|
||||
"""
|
||||
Args:
|
||||
delta: localisation
|
||||
anchor: anchor of pre-determined anchor size
|
||||
Return:
|
||||
delta: prediction of bounding box
|
||||
"""
|
||||
delta_transpose = np.transpose(delta, (1, 2, 3, 0))
|
||||
delta_contig = np.ascontiguousarray(delta_transpose)
|
||||
delta = delta_contig.reshape(4, -1)
|
||||
delta[0, :] = delta[0, :] * anchor[:, 2] + anchor[:, 0]
|
||||
delta[1, :] = delta[1, :] * anchor[:, 3] + anchor[:, 1]
|
||||
delta[2, :] = np.exp(delta[2, :]) * anchor[:, 2]
|
||||
delta[3, :] = np.exp(delta[3, :]) * anchor[:, 3]
|
||||
return delta
|
||||
|
||||
def _softmax(self, x):
|
||||
"""
|
||||
Softmax in the direction of the depth of the layer
|
||||
"""
|
||||
x = x.astype(dtype=np.float32)
|
||||
x_max = x.max(axis=1)[:, np.newaxis]
|
||||
e_x = np.exp(x-x_max)
|
||||
div = np.sum(e_x, axis=1)[:, np.newaxis]
|
||||
y = e_x / div
|
||||
return y
|
||||
|
||||
def _convert_score(self, score):
|
||||
"""
|
||||
Args:
|
||||
cls: score
|
||||
Return:
|
||||
cls: score for cls
|
||||
"""
|
||||
score_transpose = np.transpose(score, (1, 2, 3, 0))
|
||||
score_con = np.ascontiguousarray(score_transpose)
|
||||
score_view = score_con.reshape(2, -1)
|
||||
score = np.transpose(score_view, (1, 0))
|
||||
score = self._softmax(score)
|
||||
return score[:,1]
|
||||
|
||||
def _bbox_clip(self, cx, cy, width, height, boundary):
|
||||
"""
|
||||
Adjusting the bounding box
|
||||
"""
|
||||
bbox_h, bbox_w = boundary
|
||||
cx = max(0, min(cx, bbox_w))
|
||||
cy = max(0, min(cy, bbox_h))
|
||||
width = max(10, min(width, bbox_w))
|
||||
height = max(10, min(height, bbox_h))
|
||||
return cx, cy, width, height
|
||||
|
||||
def init(self, img, bbox):
|
||||
"""
|
||||
Args:
|
||||
img(np.ndarray): bgr based input image frame
|
||||
bbox: (x, y, w, h): bounding box
|
||||
"""
|
||||
x, y, w, h = bbox
|
||||
self.center_pos = np.array([x + (w - 1) / 2, y + (h - 1) / 2])
|
||||
self.h = h
|
||||
self.w = w
|
||||
w_z = self.w + self.track_context_amount * np.add(h, w)
|
||||
h_z = self.h + self.track_context_amount * np.add(h, w)
|
||||
s_z = round(np.sqrt(w_z * h_z))
|
||||
self.channel_average = np.mean(img, axis=(0, 1))
|
||||
z_crop = self.get_subwindow(img, self.center_pos, self.track_exemplar_size, s_z, self.channel_average)
|
||||
self.model.template(z_crop)
|
||||
|
||||
def track(self, img):
|
||||
"""
|
||||
Args:
|
||||
img(np.ndarray): BGR image
|
||||
Return:
|
||||
bbox(list):[x, y, width, height]
|
||||
"""
|
||||
w_z = self.w + self.track_context_amount * np.add(self.w, self.h)
|
||||
h_z = self.h + self.track_context_amount * np.add(self.w, self.h)
|
||||
s_z = np.sqrt(w_z * h_z)
|
||||
scale_z = self.track_exemplar_size / s_z
|
||||
s_x = s_z * (self.track_instance_size / self.track_exemplar_size)
|
||||
x_crop = self.get_subwindow(img, self.center_pos, self.track_instance_size, round(s_x), self.channel_average)
|
||||
outputs = self.model.track(x_crop)
|
||||
score = self._convert_score(outputs['cls'])
|
||||
pred_bbox = self._convert_bbox(outputs['loc'], self.anchors)
|
||||
|
||||
def change(r):
|
||||
return np.maximum(r, 1. / r)
|
||||
|
||||
def sz(w, h):
|
||||
pad = (w + h) * 0.5
|
||||
return np.sqrt((w + pad) * (h + pad))
|
||||
|
||||
# scale penalty
|
||||
s_c = change(sz(pred_bbox[2, :], pred_bbox[3, :]) /
|
||||
(sz(self.w * scale_z, self.h * scale_z)))
|
||||
|
||||
# aspect ratio penalty
|
||||
r_c = change((self.w / self.h) /
|
||||
(pred_bbox[2, :] / pred_bbox[3, :]))
|
||||
penalty = np.exp(-(r_c * s_c - 1) * self.track_penalty_k)
|
||||
pscore = penalty * score
|
||||
|
||||
# window penalty
|
||||
pscore = pscore * (1 - self.track_window_influence) + \
|
||||
self.window * self.track_window_influence
|
||||
best_idx = np.argmax(pscore)
|
||||
bbox = pred_bbox[:, best_idx] / scale_z
|
||||
lr = penalty[best_idx] * score[best_idx] * self.track_lr
|
||||
|
||||
cpx, cpy = self.center_pos
|
||||
x,y,w,h = bbox
|
||||
cx = x + cpx
|
||||
cy = y + cpy
|
||||
|
||||
# smooth bbox
|
||||
width = self.w * (1 - lr) + w * lr
|
||||
height = self.h * (1 - lr) + h * lr
|
||||
|
||||
# clip boundary
|
||||
cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
|
||||
|
||||
# update state
|
||||
self.center_pos = np.array([cx, cy])
|
||||
self.w = width
|
||||
self.h = height
|
||||
bbox = [cx - width / 2, cy - height / 2, width, height]
|
||||
best_score = score[best_idx]
|
||||
return {'bbox': bbox, 'best_score': best_score}
|
||||
|
||||
def get_frames(video_name):
|
||||
"""
|
||||
Args:
|
||||
Path to input video frame
|
||||
Return:
|
||||
Frame
|
||||
"""
|
||||
cap = cv.VideoCapture(video_name if video_name else 0)
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if ret:
|
||||
yield frame
|
||||
else:
|
||||
break
|
||||
|
||||
def main():
|
||||
""" Sample SiamRPN Tracker
|
||||
"""
|
||||
# Computation backends supported by layers
|
||||
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV,
|
||||
cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
|
||||
# Target Devices for computation
|
||||
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD,
|
||||
cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Use this script to run SiamRPN++ Visual Tracker',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input_video', type=str, help='Path to input video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--target_net', type=str, default='target_net.onnx', help='Path to part of SiamRPN++ ran on target frame.')
|
||||
parser.add_argument('--search_net', type=str, default='search_net.onnx', help='Path to part of SiamRPN++ ran on search frame.')
|
||||
parser.add_argument('--rpn_head', type=str, default='rpn_head.onnx', help='Path to RPN Head ONNX model.')
|
||||
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
|
||||
help="Select a computation backend: "
|
||||
"%d: automatically (by default), "
|
||||
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"%d: OpenCV Implementation, "
|
||||
"%d: VKCOM, "
|
||||
"%d: CUDA" % backends)
|
||||
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
|
||||
help='Select a target device: '
|
||||
'%d: CPU target (by default), '
|
||||
'%d: OpenCL, '
|
||||
'%d: OpenCL FP16, '
|
||||
'%d: Myriad, '
|
||||
'%d: Vulkan, '
|
||||
'%d: CUDA, '
|
||||
'%d: CUDA fp16 (half-float preprocess)' % targets)
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
if args.input_video and not os.path.isfile(args.input_video):
|
||||
raise OSError("Input video file does not exist")
|
||||
if not os.path.isfile(args.target_net):
|
||||
raise OSError("Target Net does not exist")
|
||||
if not os.path.isfile(args.search_net):
|
||||
raise OSError("Search Net does not exist")
|
||||
if not os.path.isfile(args.rpn_head):
|
||||
raise OSError("RPN Head Net does not exist")
|
||||
|
||||
#Load the Networks
|
||||
target_net = cv.dnn.readNetFromONNX(args.target_net)
|
||||
target_net.setPreferableBackend(args.backend)
|
||||
target_net.setPreferableTarget(args.target)
|
||||
search_net = cv.dnn.readNetFromONNX(args.search_net)
|
||||
search_net.setPreferableBackend(args.backend)
|
||||
search_net.setPreferableTarget(args.target)
|
||||
rpn_head = cv.dnn.readNetFromONNX(args.rpn_head)
|
||||
rpn_head.setPreferableBackend(args.backend)
|
||||
rpn_head.setPreferableTarget(args.target)
|
||||
model = ModelBuilder(target_net, search_net, rpn_head)
|
||||
tracker = SiamRPNTracker(model)
|
||||
|
||||
first_frame = True
|
||||
cv.namedWindow('SiamRPN++ Tracker', cv.WINDOW_AUTOSIZE)
|
||||
for frame in get_frames(args.input_video):
|
||||
if first_frame:
|
||||
try:
|
||||
init_rect = cv.selectROI('SiamRPN++ Tracker', frame, False, False)
|
||||
except:
|
||||
exit()
|
||||
tracker.init(frame, init_rect)
|
||||
first_frame = False
|
||||
else:
|
||||
outputs = tracker.track(frame)
|
||||
bbox = list(map(int, outputs['bbox']))
|
||||
x,y,w,h = bbox
|
||||
cv.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 3)
|
||||
cv.imshow('SiamRPN++ Tracker', frame)
|
||||
key = cv.waitKey(1)
|
||||
if key == ord("q"):
|
||||
break
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,587 @@
|
||||
#include <opencv2/core.hpp>
|
||||
#include <opencv2/videoio.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <cmath>
|
||||
#include <random>
|
||||
#include <numeric>
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
|
||||
class FilterbankFeatures {
|
||||
|
||||
// Initializes pre-processing class. Default values are the values used by the Jasper
|
||||
// architecture for pre-processing. For more details, refer to the paper here:
|
||||
// https://arxiv.org/abs/1904.03288
|
||||
|
||||
private:
|
||||
int sample_rate = 16000;
|
||||
double window_size = 0.02;
|
||||
double window_stride = 0.01;
|
||||
int win_length = static_cast<int>(sample_rate * window_size); // Number of samples in window
|
||||
int hop_length = static_cast<int>(sample_rate * window_stride); // Number of steps to advance between frames
|
||||
int n_fft = 512; // Size of window for STFT
|
||||
|
||||
// Parameters for filterbanks calculation
|
||||
int n_filt = 64;
|
||||
double lowfreq = 0.;
|
||||
double highfreq = sample_rate / 2;
|
||||
|
||||
public:
|
||||
// Mel filterbanks preparation
|
||||
double hz_to_mel(double frequencies)
|
||||
{
|
||||
//Converts frequencies from hz to mel scale
|
||||
// Fill in the linear scale
|
||||
double f_min = 0.0;
|
||||
double f_sp = 200.0 / 3;
|
||||
double mels = (frequencies - f_min) / f_sp;
|
||||
// Fill in the log-scale part
|
||||
double min_log_hz = 1000.0; // beginning of log region (Hz)
|
||||
double min_log_mel = (min_log_hz - f_min) / f_sp; // same (Mels)
|
||||
double logstep = std::log(6.4) / 27.0; // step size for log region
|
||||
|
||||
if (frequencies >= min_log_hz)
|
||||
{
|
||||
mels = min_log_mel + std::log(frequencies / min_log_hz) / logstep;
|
||||
}
|
||||
return mels;
|
||||
}
|
||||
|
||||
vector<double> mel_to_hz(vector<double>& mels)
|
||||
{
|
||||
// Converts frequencies from mel to hz scale
|
||||
|
||||
// Fill in the linear scale
|
||||
double f_min = 0.0;
|
||||
double f_sp = 200.0 / 3;
|
||||
vector<double> freqs;
|
||||
for (size_t i = 0; i < mels.size(); i++)
|
||||
{
|
||||
freqs.push_back(f_min + f_sp * mels[i]);
|
||||
}
|
||||
|
||||
// And now the nonlinear scale
|
||||
double min_log_hz = 1000.0; // beginning of log region (Hz)
|
||||
double min_log_mel = (min_log_hz - f_min) / f_sp; // same (Mels)
|
||||
double logstep = std::log(6.4) / 27.0; // step size for log region
|
||||
|
||||
for(size_t i = 0; i < mels.size(); i++)
|
||||
{
|
||||
if (mels[i] >= min_log_mel)
|
||||
{
|
||||
freqs[i] = min_log_hz * exp(logstep * (mels[i] - min_log_mel));
|
||||
}
|
||||
}
|
||||
return freqs;
|
||||
}
|
||||
|
||||
vector<double> mel_frequencies(int n_mels, double fmin, double fmax)
|
||||
{
|
||||
// Calculates n mel frequencies between 2 frequencies
|
||||
double min_mel = hz_to_mel(fmin);
|
||||
double max_mel = hz_to_mel(fmax);
|
||||
|
||||
vector<double> mels;
|
||||
double step = (max_mel - min_mel) / (n_mels - 1);
|
||||
for(double i = min_mel; i < max_mel; i += step)
|
||||
{
|
||||
mels.push_back(i);
|
||||
}
|
||||
mels.push_back(max_mel);
|
||||
|
||||
vector<double> res = mel_to_hz(mels);
|
||||
return res;
|
||||
}
|
||||
|
||||
vector<vector<double>> mel(int n_mels, double fmin, double fmax)
|
||||
{
|
||||
// Generates mel filterbank matrix
|
||||
|
||||
double num = 1 + n_fft / 2;
|
||||
vector<vector<double>> weights(n_mels, vector<double>(static_cast<int>(num), 0.));
|
||||
|
||||
// Center freqs of each FFT bin
|
||||
vector<double> fftfreqs;
|
||||
double step = (sample_rate / 2) / (num - 1);
|
||||
for(double i = 0; i <= sample_rate / 2; i += step)
|
||||
{
|
||||
fftfreqs.push_back(i);
|
||||
}
|
||||
// 'Center freqs' of mel bands - uniformly spaced between limits
|
||||
vector<double> mel_f = mel_frequencies(n_mels + 2, fmin, fmax);
|
||||
|
||||
vector<double> fdiff;
|
||||
for(size_t i = 1; i < mel_f.size(); ++i)
|
||||
{
|
||||
fdiff.push_back(mel_f[i]- mel_f[i - 1]);
|
||||
}
|
||||
|
||||
vector<vector<double>> ramps(mel_f.size(), vector<double>(fftfreqs.size()));
|
||||
for (size_t i = 0; i < mel_f.size(); ++i)
|
||||
{
|
||||
for (size_t j = 0; j < fftfreqs.size(); ++j)
|
||||
{
|
||||
ramps[i][j] = mel_f[i] - fftfreqs[j];
|
||||
}
|
||||
}
|
||||
|
||||
double lower, upper, enorm;
|
||||
for (int i = 0; i < n_mels; ++i)
|
||||
{
|
||||
// using Slaney-style mel which is scaled to be approx constant energy per channel
|
||||
enorm = 2./(mel_f[i + 2] - mel_f[i]);
|
||||
|
||||
for (int j = 0; j < static_cast<int>(num); ++j)
|
||||
{
|
||||
// lower and upper slopes for all bins
|
||||
lower = (-1) * ramps[i][j] / fdiff[i];
|
||||
upper = ramps[i + 2][j] / fdiff[i + 1];
|
||||
|
||||
weights[i][j] = max(0., min(lower, upper)) * enorm;
|
||||
}
|
||||
}
|
||||
return weights;
|
||||
}
|
||||
|
||||
// STFT preparation
|
||||
vector<double> pad_window_center(vector<double>&data, int size)
|
||||
{
|
||||
// Pad the window out to n_fft size
|
||||
int n = static_cast<int>(data.size());
|
||||
int lpad = static_cast<int>((size - n) / 2);
|
||||
vector<double> pad_array;
|
||||
|
||||
for(int i = 0; i < lpad; ++i)
|
||||
{
|
||||
pad_array.push_back(0.);
|
||||
}
|
||||
|
||||
for(size_t i = 0; i < data.size(); ++i)
|
||||
{
|
||||
pad_array.push_back(data[i]);
|
||||
}
|
||||
|
||||
for(int i = 0; i < lpad; ++i)
|
||||
{
|
||||
pad_array.push_back(0.);
|
||||
}
|
||||
return pad_array;
|
||||
}
|
||||
|
||||
vector<vector<double>> frame(vector<double>& x)
|
||||
{
|
||||
// Slices a data array into overlapping frames.
|
||||
int n_frames = static_cast<int>(1 + (x.size() - n_fft) / hop_length);
|
||||
vector<vector<double>> new_x(n_fft, vector<double>(n_frames));
|
||||
|
||||
for (int i = 0; i < n_fft; ++i)
|
||||
{
|
||||
for (int j = 0; j < n_frames; ++j)
|
||||
{
|
||||
new_x[i][j] = x[i + j * hop_length];
|
||||
}
|
||||
}
|
||||
return new_x;
|
||||
}
|
||||
|
||||
vector<double> hanning()
|
||||
{
|
||||
// https://en.wikipedia.org/wiki/Window_function#Hann_and_Hamming_windows
|
||||
vector<double> window_tensor;
|
||||
for (int j = 1 - win_length; j < win_length; j+=2)
|
||||
{
|
||||
window_tensor.push_back(1 - (0.5 * (1 - cos(CV_PI * j / (win_length - 1)))));
|
||||
}
|
||||
return window_tensor;
|
||||
}
|
||||
|
||||
vector<vector<double>> stft_power(vector<double>& y)
|
||||
{
|
||||
// Short Time Fourier Transform. The STFT represents a signal in the time-frequency
|
||||
// domain by computing discrete Fourier transforms (DFT) over short overlapping windows.
|
||||
// https://en.wikipedia.org/wiki/Short-time_Fourier_transform
|
||||
|
||||
// Pad the time series so that frames are centered
|
||||
vector<double> new_y;
|
||||
int num = int(n_fft / 2);
|
||||
|
||||
for (int i = 0; i < num; ++i)
|
||||
{
|
||||
new_y.push_back(y[num - i]);
|
||||
}
|
||||
for (size_t i = 0; i < y.size(); ++i)
|
||||
{
|
||||
new_y.push_back(y[i]);
|
||||
}
|
||||
for (size_t i = y.size() - 2; i >= y.size() - num - 1; --i)
|
||||
{
|
||||
new_y.push_back(y[i]);
|
||||
}
|
||||
|
||||
// Compute a window function
|
||||
vector<double> window_tensor = hanning();
|
||||
|
||||
// Pad the window out to n_fft size
|
||||
vector<double> fft_window = pad_window_center(window_tensor, n_fft);
|
||||
|
||||
// Window the time series
|
||||
vector<vector<double>> y_frames = frame(new_y);
|
||||
|
||||
// Multiply on fft_window
|
||||
for (size_t i = 0; i < y_frames.size(); ++i)
|
||||
{
|
||||
for (size_t j = 0; j < y_frames[0].size(); ++j)
|
||||
{
|
||||
y_frames[i][j] *= fft_window[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Transpose frames for computing stft
|
||||
vector<vector<double>> y_frames_transpose(y_frames[0].size(), vector<double>(y_frames.size()));
|
||||
for (size_t i = 0; i < y_frames[0].size(); ++i)
|
||||
{
|
||||
for (size_t j = 0; j < y_frames.size(); ++j)
|
||||
{
|
||||
y_frames_transpose[i][j] = y_frames[j][i];
|
||||
}
|
||||
}
|
||||
|
||||
// Short Time Fourier Transform
|
||||
// and get power of spectrum
|
||||
vector<vector<double>> spectrum_power(y_frames_transpose[0].size() / 2 + 1 );
|
||||
for (size_t i = 0; i < y_frames_transpose.size(); ++i)
|
||||
{
|
||||
Mat dstMat;
|
||||
dft(y_frames_transpose[i], dstMat, DFT_COMPLEX_OUTPUT);
|
||||
|
||||
// we need only the first part of the spectrum, the second part is symmetrical
|
||||
for (int j = 0; j < static_cast<int>(y_frames_transpose[0].size()) / 2 + 1; ++j)
|
||||
{
|
||||
double power_re = dstMat.at<double>(2 * j) * dstMat.at<double>(2 * j);
|
||||
double power_im = dstMat.at<double>(2 * j + 1) * dstMat.at<double>(2 * j + 1);
|
||||
spectrum_power[j].push_back(power_re + power_im);
|
||||
}
|
||||
}
|
||||
return spectrum_power;
|
||||
}
|
||||
|
||||
Mat calculate_features(vector<double>& x)
|
||||
{
|
||||
// Calculates filterbank features matrix.
|
||||
|
||||
// Do preemphasis
|
||||
std::default_random_engine generator;
|
||||
std::normal_distribution<double> normal_distr(0, 1);
|
||||
double dither = 1e-5;
|
||||
for(size_t i = 0; i < x.size(); ++i)
|
||||
{
|
||||
x[i] += dither * static_cast<double>(normal_distr(generator));
|
||||
}
|
||||
double preemph = 0.97;
|
||||
for (size_t i = x.size() - 1; i > 0; --i)
|
||||
{
|
||||
x[i] -= preemph * x[i-1];
|
||||
}
|
||||
|
||||
// Calculate Short Time Fourier Transform and get power of spectrum
|
||||
auto spectrum_power = stft_power(x);
|
||||
|
||||
vector<vector<double>> filterbanks = mel(n_filt, lowfreq, highfreq);
|
||||
|
||||
// Calculate log of multiplication of filterbanks matrix on spectrum_power matrix
|
||||
vector<vector<double>> x_stft(filterbanks.size(), vector<double>(spectrum_power[0].size(), 0));
|
||||
|
||||
for (size_t i = 0; i < filterbanks.size(); ++i)
|
||||
{
|
||||
for (size_t j = 0; j < filterbanks[0].size(); ++j)
|
||||
{
|
||||
for (size_t k = 0; k < spectrum_power[0].size(); ++k)
|
||||
{
|
||||
x_stft[i][k] += filterbanks[i][j] * spectrum_power[j][k];
|
||||
}
|
||||
}
|
||||
for (size_t k = 0; k < spectrum_power[0].size(); ++k)
|
||||
{
|
||||
x_stft[i][k] = std::log(x_stft[i][k] + 1e-20);
|
||||
}
|
||||
}
|
||||
|
||||
// normalize data
|
||||
auto elments_num = x_stft[0].size();
|
||||
for(size_t i = 0; i < x_stft.size(); ++i)
|
||||
{
|
||||
double x_mean = std::accumulate(x_stft[i].begin(), x_stft[i].end(), 0.) / elments_num; // arithmetic mean
|
||||
double x_std = 0; // standard deviation
|
||||
for(size_t j = 0; j < elments_num; ++j)
|
||||
{
|
||||
double subtract = x_stft[i][j] - x_mean;
|
||||
x_std += subtract * subtract;
|
||||
}
|
||||
x_std /= elments_num;
|
||||
x_std = sqrt(x_std) + 1e-10; // make sure x_std is not zero
|
||||
|
||||
for(size_t j = 0; j < elments_num; ++j)
|
||||
{
|
||||
x_stft[i][j] = (x_stft[i][j] - x_mean) / x_std; // standard score
|
||||
}
|
||||
}
|
||||
|
||||
Mat calculate_features(static_cast<int>(x_stft.size()), static_cast<int>(x_stft[0].size()), CV_32F);
|
||||
for(int i = 0; i < calculate_features.size[0]; ++i)
|
||||
{
|
||||
for(int j = 0; j < calculate_features.size[1]; ++j)
|
||||
{
|
||||
calculate_features.at<float>(i, j) = static_cast<float>(x_stft[i][j]);
|
||||
}
|
||||
}
|
||||
return calculate_features;
|
||||
}
|
||||
};
|
||||
|
||||
class Decoder {
|
||||
// Used for decoding the output of jasper model
|
||||
private:
|
||||
unordered_map<int, char> labels_map = fillMap();
|
||||
int blank_id = 28;
|
||||
|
||||
public:
|
||||
unordered_map<int, char> fillMap()
|
||||
{
|
||||
vector<char> labels={' ','a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p'
|
||||
,'q','r','s','t','u','v','w','x','y','z','\''};
|
||||
unordered_map<int, char> map;
|
||||
for(int i = 0; i < static_cast<int>(labels.size()); ++i)
|
||||
{
|
||||
map[i] = labels[i];
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
string decode(Mat& x)
|
||||
{
|
||||
// Takes output of Jasper model and performs ctc decoding algorithm to
|
||||
// remove duplicates and special symbol. Returns prediction
|
||||
|
||||
vector<int> prediction;
|
||||
for(int i = 0; i < x.size[1]; ++i)
|
||||
{
|
||||
double maxEl = -1e10;
|
||||
int ind = 0;
|
||||
for(int j = 0; j < x.size[2]; ++j)
|
||||
{
|
||||
if (maxEl <= x.at<float>(0, i, j))
|
||||
{
|
||||
maxEl = x.at<float>(0, i, j);
|
||||
ind = j;
|
||||
}
|
||||
}
|
||||
prediction.push_back(ind);
|
||||
}
|
||||
// CTC decoding procedure
|
||||
vector<double> decoded_prediction = {};
|
||||
int previous = blank_id;
|
||||
|
||||
for(int i = 0; i < static_cast<int>(prediction.size()); ++i)
|
||||
{
|
||||
if (( prediction[i] != previous || previous == blank_id) && prediction[i] != blank_id)
|
||||
{
|
||||
decoded_prediction.push_back(prediction[i]);
|
||||
}
|
||||
previous = prediction[i];
|
||||
}
|
||||
|
||||
string hypotheses = {};
|
||||
for(size_t i = 0; i < decoded_prediction.size(); ++i)
|
||||
{
|
||||
auto it = labels_map.find(static_cast<char>(decoded_prediction[i]));
|
||||
if (it != labels_map.end())
|
||||
hypotheses.push_back(it->second);
|
||||
}
|
||||
return hypotheses;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
static string predict(Mat& features, dnn::Net net, Decoder decoder)
|
||||
{
|
||||
// Passes the features through the Jasper model and decodes the output to english transcripts.
|
||||
|
||||
// expand 2d features matrix to 3d
|
||||
vector<int> sizes = {1, static_cast<int>(features.size[0]),
|
||||
static_cast<int>(features.size[1])};
|
||||
features = features.reshape(0, sizes);
|
||||
|
||||
// make prediction
|
||||
net.setInput(features);
|
||||
Mat output = net.forward();
|
||||
|
||||
// decode output to transcript
|
||||
auto prediction = decoder.decode(output);
|
||||
return prediction;
|
||||
}
|
||||
|
||||
static int readAudioFile(vector<double>& inputAudio, string file, int audioStream)
|
||||
{
|
||||
VideoCapture cap;
|
||||
int samplingRate = 16000;
|
||||
vector<int> params { CAP_PROP_AUDIO_STREAM, audioStream,
|
||||
CAP_PROP_VIDEO_STREAM, -1,
|
||||
CAP_PROP_AUDIO_DATA_DEPTH, CV_32F,
|
||||
CAP_PROP_AUDIO_SAMPLES_PER_SECOND, samplingRate
|
||||
};
|
||||
cap.open(file, CAP_ANY, params);
|
||||
if (!cap.isOpened())
|
||||
{
|
||||
cerr << "Error : Can't read audio file: '" << file << "' with audioStream = " << audioStream << endl;
|
||||
return -1;
|
||||
}
|
||||
const int audioBaseIndex = (int)cap.get(CAP_PROP_AUDIO_BASE_INDEX);
|
||||
vector<double> frameVec;
|
||||
Mat frame;
|
||||
for (;;)
|
||||
{
|
||||
if (cap.grab())
|
||||
{
|
||||
cap.retrieve(frame, audioBaseIndex);
|
||||
frameVec = frame;
|
||||
inputAudio.insert(inputAudio.end(), frameVec.begin(), frameVec.end());
|
||||
}
|
||||
else
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
return samplingRate;
|
||||
}
|
||||
|
||||
static int readAudioMicrophone(vector<double>& inputAudio, int microTime)
|
||||
{
|
||||
VideoCapture cap;
|
||||
int samplingRate = 16000;
|
||||
vector<int> params { CAP_PROP_AUDIO_STREAM, 0,
|
||||
CAP_PROP_VIDEO_STREAM, -1,
|
||||
CAP_PROP_AUDIO_DATA_DEPTH, CV_32F,
|
||||
CAP_PROP_AUDIO_SAMPLES_PER_SECOND, samplingRate
|
||||
};
|
||||
cap.open(0, CAP_ANY, params);
|
||||
if (!cap.isOpened())
|
||||
{
|
||||
cerr << "Error: Can't open microphone" << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
const int audioBaseIndex = (int)cap.get(CAP_PROP_AUDIO_BASE_INDEX);
|
||||
vector<double> frameVec;
|
||||
Mat frame;
|
||||
if (microTime <= 0)
|
||||
{
|
||||
cerr << "Error: Duration of audio chunk must be > 0" << endl;
|
||||
return -1;
|
||||
}
|
||||
size_t sizeOfData = static_cast<size_t>(microTime * samplingRate);
|
||||
while (inputAudio.size() < sizeOfData)
|
||||
{
|
||||
if (cap.grab())
|
||||
{
|
||||
cap.retrieve(frame, audioBaseIndex);
|
||||
frameVec = frame;
|
||||
inputAudio.insert(inputAudio.end(), frameVec.begin(), frameVec.end());
|
||||
}
|
||||
else
|
||||
{
|
||||
cerr << "Error: Grab error" << endl;
|
||||
break;
|
||||
}
|
||||
}
|
||||
return samplingRate;
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
const String keys =
|
||||
"{help h usage ? | | This script runs Jasper Speech recognition model }"
|
||||
"{input_file i | | Path to input audio file. If not specified, microphone input will be used }"
|
||||
"{audio_duration t | 15 | Duration of audio chunk to be captured from microphone }"
|
||||
"{audio_stream a | 0 | CAP_PROP_AUDIO_STREAM value }"
|
||||
"{show_spectrogram s | false | Show a spectrogram of the input audio: true / false / 1 / 0 }"
|
||||
"{model m | jasper.onnx | Path to the onnx file of Jasper. You can download the converted onnx model "
|
||||
"from https://drive.google.com/drive/folders/1wLtxyao4ItAg8tt4Sb63zt6qXzhcQoR6?usp=sharing}"
|
||||
"{backend b | dnn::DNN_BACKEND_DEFAULT | Select a computation backend: "
|
||||
"dnn::DNN_BACKEND_DEFAULT, "
|
||||
"dnn::DNN_BACKEND_INFERENCE_ENGINE, "
|
||||
"dnn::DNN_BACKEND_OPENCV }"
|
||||
"{target t | dnn::DNN_TARGET_CPU | Select a target device: "
|
||||
"dnn::DNN_TARGET_CPU, "
|
||||
"dnn::DNN_TARGET_OPENCL, "
|
||||
"dnn::DNN_TARGET_OPENCL_FP16 }"
|
||||
;
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
if (parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Load Network
|
||||
dnn::Net net = dnn::readNetFromONNX(parser.get<std::string>("model"));
|
||||
net.setPreferableBackend(parser.get<int>("backend"));
|
||||
net.setPreferableTarget(parser.get<int>("target"));
|
||||
|
||||
// Get audio
|
||||
vector<double>inputAudio = {};
|
||||
int samplingRate = 0;
|
||||
if (parser.has("input_file"))
|
||||
{
|
||||
string audio = samples::findFile(parser.get<std::string>("input_file"));
|
||||
samplingRate = readAudioFile(inputAudio, audio, parser.get<int>("audio_stream"));
|
||||
}
|
||||
else
|
||||
{
|
||||
samplingRate = readAudioMicrophone(inputAudio, parser.get<int>("audio_duration"));
|
||||
}
|
||||
|
||||
if ((inputAudio.size() == 0) || samplingRate <= 0)
|
||||
{
|
||||
cerr << "Error: problems with audio reading, check input arguments" << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (inputAudio.size() / samplingRate < 6)
|
||||
{
|
||||
cout << "Warning: For predictable network performance duration of audio must exceed 6 sec."
|
||||
" Audio will be extended with zero samples" << endl;
|
||||
for(int i = static_cast<int>(inputAudio.size()) - 1; i < samplingRate * 6; ++i)
|
||||
{
|
||||
inputAudio.push_back(0);
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate features
|
||||
FilterbankFeatures filter;
|
||||
auto calculated_features = filter.calculate_features(inputAudio);
|
||||
|
||||
// Show spectogram if required
|
||||
if (parser.get<bool>("show_spectrogram") == true)
|
||||
{
|
||||
Mat spectogram;
|
||||
normalize(calculated_features, spectogram, 0, 255, NORM_MINMAX, CV_8U);
|
||||
applyColorMap(spectogram, spectogram, COLORMAP_INFERNO);
|
||||
imshow("spectogram", spectogram);
|
||||
waitKey(0);
|
||||
}
|
||||
|
||||
Decoder decoder;
|
||||
string prediction = predict(calculated_features, net, decoder);
|
||||
for( auto &transcript: prediction)
|
||||
{
|
||||
cout << transcript;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,567 @@
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import os
|
||||
|
||||
'''
|
||||
You can download the converted onnx model from https://drive.google.com/drive/folders/1wLtxyao4ItAg8tt4Sb63zt6qXzhcQoR6?usp=sharing
|
||||
or convert the model yourself.
|
||||
|
||||
You can get the original pre-trained Jasper model from NVIDIA : https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_onnx_fp16_amp/files
|
||||
Download and unzip : `$ wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/jasper_pyt_onnx_fp16_amp/versions/20.10.0/zip -O jasper_pyt_onnx_fp16_amp_20.10.0.zip && unzip -o ./jasper_pyt_onnx_fp16_amp_20.10.0.zip && unzip -o ./jasper_pyt_onnx_fp16_amp.zip`
|
||||
|
||||
you can get the script to convert the model here : https://gist.github.com/spazewalker/507f1529e19aea7e8417f6e935851a01
|
||||
|
||||
You can convert the model using the following steps:
|
||||
1. Import onnx and load the original model
|
||||
```
|
||||
import onnx
|
||||
model = onnx.load("./jasper-onnx/1/model.onnx")
|
||||
```
|
||||
|
||||
3. Change data type of input layer
|
||||
```
|
||||
inp = model.graph.input[0]
|
||||
model.graph.input.remove(inp)
|
||||
inp.type.tensor_type.elem_type = 1
|
||||
model.graph.input.insert(0,inp)
|
||||
```
|
||||
|
||||
4. Change the data type of output layer
|
||||
```
|
||||
out = model.graph.output[0]
|
||||
model.graph.output.remove(out)
|
||||
out.type.tensor_type.elem_type = 1
|
||||
model.graph.output.insert(0,out)
|
||||
```
|
||||
|
||||
5. Change the data type of every initializer and cast it's values from FP16 to FP32
|
||||
```
|
||||
for i,init in enumerate(model.graph.initializer):
|
||||
model.graph.initializer.remove(init)
|
||||
init.data_type = 1
|
||||
init.raw_data = np.frombuffer(init.raw_data, count=np.product(init.dims), dtype=np.float16).astype(np.float32).tobytes()
|
||||
model.graph.initializer.insert(i,init)
|
||||
```
|
||||
|
||||
6. Add an additional reshape node to handle the inconsistent input from python and c++ of openCV.
|
||||
see https://github.com/opencv/opencv/issues/19091
|
||||
Make & insert a new node with 'Reshape' operation & required initializer
|
||||
```
|
||||
tensor = numpy_helper.from_array(np.array([0,64,-1]),name='shape_reshape')
|
||||
model.graph.initializer.insert(0,tensor)
|
||||
node = onnx.helper.make_node(op_type='Reshape',inputs=['input__0','shape_reshape'], outputs=['input_reshaped'], name='reshape__0')
|
||||
model.graph.node.insert(0,node)
|
||||
model.graph.node[1].input[0] = 'input_reshaped'
|
||||
```
|
||||
|
||||
7. Finally save the model
|
||||
```
|
||||
with open('jasper_dynamic_input_float.onnx','wb') as f:
|
||||
onnx.save_model(model,f)
|
||||
```
|
||||
|
||||
Original Repo : https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/Jasper
|
||||
'''
|
||||
|
||||
class FilterbankFeatures:
|
||||
def __init__(self,
|
||||
sample_rate=16000, window_size=0.02, window_stride=0.01,
|
||||
n_fft=512, preemph=0.97, n_filt=64, lowfreq=0,
|
||||
highfreq=None, log=True, dither=1e-5):
|
||||
'''
|
||||
Initializes pre-processing class. Default values are the values used by the Jasper
|
||||
architecture for pre-processing. For more details, refer to the paper here:
|
||||
https://arxiv.org/abs/1904.03288
|
||||
'''
|
||||
self.win_length = int(sample_rate * window_size) # frame size
|
||||
self.hop_length = int(sample_rate * window_stride) # stride
|
||||
self.n_fft = n_fft or 2 ** np.ceil(np.log2(self.win_length))
|
||||
self.log = log
|
||||
self.dither = dither
|
||||
self.n_filt = n_filt
|
||||
self.preemph = preemph
|
||||
highfreq = highfreq or sample_rate / 2
|
||||
self.window_tensor = np.hanning(self.win_length)
|
||||
|
||||
self.filterbanks = self.mel(sample_rate, self.n_fft, n_mels=n_filt, fmin=lowfreq, fmax=highfreq)
|
||||
self.filterbanks.dtype=np.float32
|
||||
self.filterbanks = np.expand_dims(self.filterbanks,0)
|
||||
|
||||
def normalize_batch(self, x, seq_len):
|
||||
'''
|
||||
Normalizes the features.
|
||||
'''
|
||||
x_mean = np.zeros((seq_len.shape[0], x.shape[1]), dtype=x.dtype)
|
||||
x_std = np.zeros((seq_len.shape[0], x.shape[1]), dtype=x.dtype)
|
||||
for i in range(x.shape[0]):
|
||||
x_mean[i, :] = np.mean(x[i, :, :seq_len[i]],axis=1)
|
||||
x_std[i, :] = np.std(x[i, :, :seq_len[i]],axis=1)
|
||||
# make sure x_std is not zero
|
||||
x_std += 1e-10
|
||||
return (x - np.expand_dims(x_mean,2)) / np.expand_dims(x_std,2)
|
||||
|
||||
def calculate_features(self, x, seq_len):
|
||||
'''
|
||||
Calculates filterbank features.
|
||||
args:
|
||||
x : mono channel audio
|
||||
seq_len : length of the audio sample
|
||||
returns:
|
||||
x : filterbank features
|
||||
'''
|
||||
dtype = x.dtype
|
||||
|
||||
seq_len = np.ceil(seq_len / self.hop_length)
|
||||
seq_len = np.array(seq_len,dtype=np.int32)
|
||||
|
||||
# dither
|
||||
if self.dither > 0:
|
||||
x += self.dither * np.random.randn(*x.shape)
|
||||
|
||||
# do preemphasis
|
||||
if self.preemph is not None:
|
||||
x = np.concatenate(
|
||||
(np.expand_dims(x[0],-1), x[1:] - self.preemph * x[:-1]), axis=0)
|
||||
|
||||
# Short Time Fourier Transform
|
||||
x = self.stft(x, n_fft=self.n_fft, hop_length=self.hop_length,
|
||||
win_length=self.win_length,
|
||||
fft_window=self.window_tensor)
|
||||
|
||||
# get power spectrum
|
||||
x = (x**2).sum(-1)
|
||||
|
||||
# dot with filterbank energies
|
||||
x = np.matmul(np.array(self.filterbanks,dtype=x.dtype), x)
|
||||
|
||||
# log features if required
|
||||
if self.log:
|
||||
x = np.log(x + 1e-20)
|
||||
|
||||
# normalize if required
|
||||
x = self.normalize_batch(x, seq_len).astype(dtype)
|
||||
return x
|
||||
|
||||
# Mel Frequency calculation
|
||||
def hz_to_mel(self, frequencies):
|
||||
'''
|
||||
Converts frequencies from hz to mel scale. Input can be a number or a vector.
|
||||
'''
|
||||
frequencies = np.asanyarray(frequencies)
|
||||
|
||||
f_min = 0.0
|
||||
f_sp = 200.0 / 3
|
||||
|
||||
mels = (frequencies - f_min) / f_sp
|
||||
|
||||
# Fill in the log-scale part
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels)
|
||||
logstep = np.log(6.4) / 27.0 # step size for log region
|
||||
|
||||
if frequencies.ndim:
|
||||
# If we have array data, vectorize
|
||||
log_t = frequencies >= min_log_hz
|
||||
mels[log_t] = min_log_mel + np.log(frequencies[log_t] / min_log_hz) / logstep
|
||||
elif frequencies >= min_log_hz:
|
||||
# If we have scalar data, directly
|
||||
mels = min_log_mel + np.log(frequencies / min_log_hz) / logstep
|
||||
return mels
|
||||
|
||||
def mel_to_hz(self, mels):
|
||||
'''
|
||||
Converts frequencies from mel to hz scale. Input can be a number or a vector.
|
||||
'''
|
||||
mels = np.asanyarray(mels)
|
||||
|
||||
# Fill in the linear scale
|
||||
f_min = 0.0
|
||||
f_sp = 200.0 / 3
|
||||
freqs = f_min + f_sp * mels
|
||||
|
||||
# And now the nonlinear scale
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels)
|
||||
logstep = np.log(6.4) / 27.0 # step size for log region
|
||||
|
||||
if mels.ndim:
|
||||
# If we have vector data, vectorize
|
||||
log_t = mels >= min_log_mel
|
||||
freqs[log_t] = min_log_hz * np.exp(logstep * (mels[log_t] - min_log_mel))
|
||||
elif mels >= min_log_mel:
|
||||
# If we have scalar data, check directly
|
||||
freqs = min_log_hz * np.exp(logstep * (mels - min_log_mel))
|
||||
|
||||
return freqs
|
||||
|
||||
def mel_frequencies(self, n_mels=128, fmin=0.0, fmax=11025.0):
|
||||
'''
|
||||
Calculates n mel frequencies between 2 frequencies
|
||||
args:
|
||||
n_mels : number of bands
|
||||
fmin : min frequency
|
||||
fmax : max frequency
|
||||
returns:
|
||||
mels : vector of mel frequencies
|
||||
'''
|
||||
# 'Center freqs' of mel bands - uniformly spaced between limits
|
||||
min_mel = self.hz_to_mel(fmin)
|
||||
max_mel = self.hz_to_mel(fmax)
|
||||
|
||||
mels = np.linspace(min_mel, max_mel, n_mels)
|
||||
|
||||
return self.mel_to_hz(mels)
|
||||
|
||||
def mel(self, sr, n_fft, n_mels=128, fmin=0.0, fmax=None, dtype=np.float32):
|
||||
'''
|
||||
Generates mel filterbank
|
||||
args:
|
||||
sr : Sampling rate
|
||||
n_fft : number of FFT components
|
||||
n_mels : number of Mel bands to generate
|
||||
fmin : lowest frequency (in Hz)
|
||||
fmax : highest frequency (in Hz). sr/2.0 if None
|
||||
dtype : the data type of the output basis.
|
||||
returns:
|
||||
mels : Mel transform matrix
|
||||
'''
|
||||
# default Max freq = half of sampling rate
|
||||
if fmax is None:
|
||||
fmax = float(sr) / 2
|
||||
|
||||
# Initialize the weights
|
||||
n_mels = int(n_mels)
|
||||
weights = np.zeros((n_mels, int(1 + n_fft // 2)), dtype=dtype)
|
||||
|
||||
# Center freqs of each FFT bin
|
||||
fftfreqs = np.linspace(0, float(sr) / 2, int(1 + n_fft // 2), endpoint=True)
|
||||
|
||||
# 'Center freqs' of mel bands - uniformly spaced between limits
|
||||
mel_f = self.mel_frequencies(n_mels + 2, fmin=fmin, fmax=fmax)
|
||||
|
||||
fdiff = np.diff(mel_f)
|
||||
ramps = np.subtract.outer(mel_f, fftfreqs)
|
||||
|
||||
for i in range(n_mels):
|
||||
# lower and upper slopes for all bins
|
||||
lower = -ramps[i] / fdiff[i]
|
||||
upper = ramps[i + 2] / fdiff[i + 1]
|
||||
|
||||
# .. then intersect them with each other and zero
|
||||
weights[i] = np.maximum(0, np.minimum(lower, upper))
|
||||
|
||||
# Using Slaney-style mel which is scaled to be approx constant energy per channel
|
||||
enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
|
||||
weights *= enorm[:, np.newaxis]
|
||||
return weights
|
||||
|
||||
# STFT preparation
|
||||
def pad_window_center(self, data, size, axis=-1, **kwargs):
|
||||
'''
|
||||
Centers the data and pads.
|
||||
args:
|
||||
data : Vector to be padded and centered
|
||||
size : Length to pad data
|
||||
axis : Axis along which to pad and center the data
|
||||
kwargs : arguments passed to np.pad
|
||||
return : centered and padded data
|
||||
'''
|
||||
kwargs.setdefault("mode", "constant")
|
||||
n = data.shape[axis]
|
||||
lpad = int((size - n) // 2)
|
||||
lengths = [(0, 0)] * data.ndim
|
||||
lengths[axis] = (lpad, int(size - n - lpad))
|
||||
if lpad < 0:
|
||||
raise Exception(
|
||||
("Target size ({:d}) must be at least input size ({:d})").format(size, n)
|
||||
)
|
||||
return np.pad(data, lengths, **kwargs)
|
||||
|
||||
def frame(self, x, frame_length, hop_length):
|
||||
'''
|
||||
Slices a data array into (overlapping) frames.
|
||||
args:
|
||||
x : array to frame
|
||||
frame_length : length of frame
|
||||
hop_length : Number of steps to advance between frames
|
||||
return : A framed view of `x`
|
||||
'''
|
||||
if x.shape[-1] < frame_length:
|
||||
raise Exception(
|
||||
"Input is too short (n={:d})"
|
||||
" for frame_length={:d}".format(x.shape[-1], frame_length)
|
||||
)
|
||||
x = np.asfortranarray(x)
|
||||
n_frames = 1 + (x.shape[-1] - frame_length) // hop_length
|
||||
strides = np.asarray(x.strides)
|
||||
new_stride = np.prod(strides[strides > 0] // x.itemsize) * x.itemsize
|
||||
shape = list(x.shape)[:-1] + [frame_length, n_frames]
|
||||
strides = list(strides) + [hop_length * new_stride]
|
||||
return np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides)
|
||||
|
||||
def dtype_r2c(self, d, default=np.complex64):
|
||||
'''
|
||||
Find the complex numpy dtype corresponding to a real dtype.
|
||||
args:
|
||||
d : The real-valued dtype to convert to complex.
|
||||
default : The default complex target type, if `d` does not match a known dtype
|
||||
return : The complex dtype
|
||||
'''
|
||||
mapping = {
|
||||
np.dtype(np.float32): np.complex64,
|
||||
np.dtype(np.float64): np.complex128,
|
||||
}
|
||||
dt = np.dtype(d)
|
||||
if dt.kind == "c":
|
||||
return dt
|
||||
return np.dtype(mapping.get(dt, default))
|
||||
|
||||
def stft(self, y, n_fft, hop_length=None, win_length=None, fft_window=None, pad_mode='reflect', return_complex=False):
|
||||
'''
|
||||
Short Time Fourier Transform. The STFT represents a signal in the time-frequency
|
||||
domain by computing discrete Fourier transforms (DFT) over short overlapping windows.
|
||||
args:
|
||||
y : input signal
|
||||
n_fft : length of the windowed signal after padding with zeros.
|
||||
hop_length : number of audio samples between adjacent STFT columns.
|
||||
win_length : Each frame of audio is windowed by window of length win_length and
|
||||
then padded with zeros to match n_fft
|
||||
fft_window : a vector or array of length `n_fft` having values computed by a
|
||||
window function
|
||||
pad_mode : mode while padding the signal
|
||||
return_complex : returns array with complex data type if `True`
|
||||
return : Matrix of short-term Fourier transform coefficients.
|
||||
'''
|
||||
if win_length is None:
|
||||
win_length = n_fft
|
||||
if hop_length is None:
|
||||
hop_length = int(win_length // 4)
|
||||
if y.ndim!=1:
|
||||
raise Exception(f'Invalid input shape. Only Mono Channeled audio supported. Input must have shape (Audio,). Got {y.shape}')
|
||||
|
||||
# Pad the window out to n_fft size
|
||||
fft_window = self.pad_window_center(fft_window, n_fft)
|
||||
|
||||
# Reshape so that the window can be broadcast
|
||||
fft_window = fft_window.reshape((-1, 1))
|
||||
|
||||
# Pad the time series so that frames are centered
|
||||
y = np.pad(y, int(n_fft // 2), mode=pad_mode)
|
||||
|
||||
# Window the time series.
|
||||
y_frames = self.frame(y, frame_length=n_fft, hop_length=hop_length)
|
||||
|
||||
# Convert data type to complex
|
||||
dtype = self.dtype_r2c(y.dtype)
|
||||
|
||||
# Pre-allocate the STFT matrix
|
||||
stft_matrix = np.empty( (int(1 + n_fft // 2), y_frames.shape[-1]), dtype=dtype, order="F")
|
||||
|
||||
stft_matrix = np.fft.rfft( fft_window * y_frames, axis=0)
|
||||
return stft_matrix if return_complex==True else np.stack((stft_matrix.real,stft_matrix.imag),axis=-1)
|
||||
|
||||
class Decoder:
|
||||
'''
|
||||
Used for decoding the output of jasper model.
|
||||
'''
|
||||
def __init__(self):
|
||||
labels=[' ','a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z',"'"]
|
||||
self.labels_map = {i: label for i,label in enumerate(labels)}
|
||||
self.blank_id = 28
|
||||
|
||||
def decode(self,x):
|
||||
"""
|
||||
Takes output of Jasper model and performs ctc decoding algorithm to
|
||||
remove duplicates and special symbol. Returns prediction
|
||||
"""
|
||||
x = np.argmax(x,axis=-1)
|
||||
hypotheses = []
|
||||
prediction = x.tolist()
|
||||
# CTC decoding procedure
|
||||
decoded_prediction = []
|
||||
previous = self.blank_id
|
||||
for p in prediction:
|
||||
if (p != previous or previous == self.blank_id) and p != self.blank_id:
|
||||
decoded_prediction.append(p)
|
||||
previous = p
|
||||
hypothesis = ''.join([self.labels_map[c] for c in decoded_prediction])
|
||||
hypotheses.append(hypothesis)
|
||||
return hypotheses
|
||||
|
||||
def predict(features, net, decoder):
|
||||
'''
|
||||
Passes the features through the Jasper model and decodes the output to english transcripts.
|
||||
args:
|
||||
features : input features, calculated using FilterbankFeatures class
|
||||
net : Jasper model dnn.net object
|
||||
decoder : Decoder object
|
||||
return : Predicted text
|
||||
'''
|
||||
# make prediction
|
||||
net.setInput(features)
|
||||
output = net.forward()
|
||||
|
||||
# decode output to transcript
|
||||
prediction = decoder.decode(output.squeeze(0))
|
||||
return prediction[0]
|
||||
|
||||
def readAudioFile(file, audioStream):
|
||||
cap = cv.VideoCapture(file)
|
||||
samplingRate = 16000
|
||||
params = np.asarray([cv.CAP_PROP_AUDIO_STREAM, audioStream,
|
||||
cv.CAP_PROP_VIDEO_STREAM, -1,
|
||||
cv.CAP_PROP_AUDIO_DATA_DEPTH, cv.CV_32F,
|
||||
cv.CAP_PROP_AUDIO_SAMPLES_PER_SECOND, samplingRate
|
||||
])
|
||||
cap.open(file, cv.CAP_ANY, params)
|
||||
if cap.isOpened() is False:
|
||||
print("Error : Can't read audio file:", file, "with audioStream = ", audioStream)
|
||||
return
|
||||
audioBaseIndex = int (cap.get(cv.CAP_PROP_AUDIO_BASE_INDEX))
|
||||
inputAudio = []
|
||||
while(1):
|
||||
if (cap.grab()):
|
||||
frame = np.asarray([])
|
||||
frame = cap.retrieve(frame, audioBaseIndex)
|
||||
for i in range(len(frame[1][0])):
|
||||
inputAudio.append(frame[1][0][i])
|
||||
else:
|
||||
break
|
||||
inputAudio = np.asarray(inputAudio, dtype=np.float64)
|
||||
return inputAudio, samplingRate
|
||||
|
||||
def readAudioMicrophone(microTime):
|
||||
cap = cv.VideoCapture()
|
||||
samplingRate = 16000
|
||||
params = np.asarray([cv.CAP_PROP_AUDIO_STREAM, 0,
|
||||
cv.CAP_PROP_VIDEO_STREAM, -1,
|
||||
cv.CAP_PROP_AUDIO_DATA_DEPTH, cv.CV_32F,
|
||||
cv.CAP_PROP_AUDIO_SAMPLES_PER_SECOND, samplingRate
|
||||
])
|
||||
cap.open(0, cv.CAP_ANY, params)
|
||||
if cap.isOpened() is False:
|
||||
print("Error: Can't open microphone")
|
||||
print("Error: problems with audio reading, check input arguments")
|
||||
return
|
||||
audioBaseIndex = int(cap.get(cv.CAP_PROP_AUDIO_BASE_INDEX))
|
||||
cvTickFreq = cv.getTickFrequency()
|
||||
sysTimeCurr = cv.getTickCount()
|
||||
sysTimePrev = sysTimeCurr
|
||||
inputAudio = []
|
||||
while ((sysTimeCurr - sysTimePrev) / cvTickFreq < microTime):
|
||||
if (cap.grab()):
|
||||
frame = np.asarray([])
|
||||
frame = cap.retrieve(frame, audioBaseIndex)
|
||||
for i in range(len(frame[1][0])):
|
||||
inputAudio.append(frame[1][0][i])
|
||||
sysTimeCurr = cv.getTickCount()
|
||||
else:
|
||||
print("Error: Grab error")
|
||||
break
|
||||
inputAudio = np.asarray(inputAudio, dtype=np.float64)
|
||||
print("Number of samples: ", len(inputAudio))
|
||||
return inputAudio, samplingRate
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
# Computation backends supported by layers
|
||||
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV)
|
||||
# Target Devices for computation
|
||||
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16)
|
||||
|
||||
parser = argparse.ArgumentParser(description='This script runs Jasper Speech recognition model',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input_type', type=str, required=True, help='file or microphone')
|
||||
parser.add_argument('--micro_time', type=int, default=15, help='Duration of microphone work in seconds. Must be more than 6 sec')
|
||||
parser.add_argument('--input_audio', type=str, help='Path to input audio file. OR Path to a txt file with relative path to multiple audio files in different lines')
|
||||
parser.add_argument('--audio_stream', type=int, default=0, help='CAP_PROP_AUDIO_STREAM value')
|
||||
parser.add_argument('--show_spectrogram', action='store_true', help='Whether to show a spectrogram of the input audio.')
|
||||
parser.add_argument('--model', type=str, default='jasper.onnx', help='Path to the onnx file of Jasper. default="jasper.onnx"')
|
||||
parser.add_argument('--output', type=str, help='Path to file where recognized audio transcript must be saved. Leave this to print on console.')
|
||||
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
|
||||
help='Select a computation backend: '
|
||||
"%d: automatically (by default) "
|
||||
"%d: OpenVINO Inference Engine "
|
||||
"%d: OpenCV Implementation " % backends)
|
||||
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
|
||||
help='Select a target device: '
|
||||
"%d: CPU target (by default) "
|
||||
"%d: OpenCL "
|
||||
"%d: OpenCL FP16 " % targets)
|
||||
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
if args.input_audio and not os.path.isfile(args.input_audio):
|
||||
raise OSError("Input audio file does not exist")
|
||||
if not os.path.isfile(args.model):
|
||||
raise OSError("Jasper model file does not exist")
|
||||
|
||||
features = []
|
||||
if args.input_type == "file":
|
||||
if args.input_audio.endswith('.txt'):
|
||||
with open(args.input_audio) as f:
|
||||
content = f.readlines()
|
||||
content = [x.strip() for x in content]
|
||||
audio_file_paths = content
|
||||
for audio_file_path in audio_file_paths:
|
||||
if not os.path.isfile(audio_file_path):
|
||||
raise OSError("Audio file({audio_file_path}) does not exist")
|
||||
else:
|
||||
audio_file_paths = [args.input_audio]
|
||||
audio_file_paths = [os.path.abspath(x) for x in audio_file_paths]
|
||||
|
||||
# Read audio Files
|
||||
for audio_file_path in audio_file_paths:
|
||||
audio = readAudioFile(audio_file_path, args.audio_stream)
|
||||
if audio is None:
|
||||
raise Exception(f"Can't read {args.input_audio}. Try a different format")
|
||||
features.append(audio[0])
|
||||
elif args.input_type == "microphone":
|
||||
# Read audio from microphone
|
||||
audio = readAudioMicrophone(args.micro_time)
|
||||
if audio is None:
|
||||
raise Exception(f"Can't open microphone. Try a different format")
|
||||
features.append(audio[0])
|
||||
else:
|
||||
raise Exception(f"input_type {args.input_type} doesn't exist. Please enter 'file' or 'microphone'")
|
||||
|
||||
# Get Filterbank Features
|
||||
feature_extractor = FilterbankFeatures()
|
||||
for i in range(len(features)):
|
||||
X = features[i]
|
||||
seq_len = np.array([X.shape[0]], dtype=np.int32)
|
||||
features[i] = feature_extractor.calculate_features(x=X, seq_len=seq_len)
|
||||
|
||||
# Load Network
|
||||
net = cv.dnn.readNetFromONNX(args.model)
|
||||
net.setPreferableBackend(args.backend)
|
||||
net.setPreferableTarget(args.target)
|
||||
|
||||
# Show spectogram if required
|
||||
if args.show_spectrogram and not args.input_audio.endswith('.txt'):
|
||||
img = cv.normalize(src=features[0][0], dst=None, alpha=0, beta=255, norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U)
|
||||
img = cv.applyColorMap(img, cv.COLORMAP_JET)
|
||||
cv.imshow('spectogram', img)
|
||||
cv.waitKey(0)
|
||||
|
||||
# Initialize decoder
|
||||
decoder = Decoder()
|
||||
|
||||
# Make prediction
|
||||
prediction = []
|
||||
print("Predicting...")
|
||||
for feature in features:
|
||||
print(f"\rAudio file {len(prediction)+1}/{len(features)}", end='')
|
||||
prediction.append(predict(feature, net, decoder))
|
||||
print("")
|
||||
|
||||
# save transcript if required
|
||||
if args.output:
|
||||
with open(args.output,'w') as f:
|
||||
for pred in prediction:
|
||||
f.write(pred+'\n')
|
||||
print("Transcript was written to {}".format(args.output))
|
||||
else:
|
||||
print(prediction)
|
||||
cv.destroyAllWindows()
|
||||
@@ -0,0 +1,198 @@
|
||||
/*
|
||||
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.
|
||||
|
||||
|
||||
This sample demonstrates super-resolution using the SeeMoreDetails model.
|
||||
The model upscales images by 4x while enhancing details and reducing noise.
|
||||
Supports image inputs only.
|
||||
|
||||
SeeMoreDetails Repo: https://github.com/eduardzamfir/seemoredetails
|
||||
*/
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <iostream>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace cv::dnn;
|
||||
using namespace std;
|
||||
|
||||
const int WINDOW_OFFSET_X = 50;
|
||||
const int WINDOW_OFFSET_Y = 50;
|
||||
const int WINDOW_SPACING = 50;
|
||||
|
||||
const string param_keys =
|
||||
"{ help h | | Print help message }"
|
||||
"{ @alias | seemoredetails | Model alias from models.yml }"
|
||||
"{ zoo | ../dnn/models.yml | Path to models.yml file }"
|
||||
"{ input i | chicky_512.png | Path to input image }"
|
||||
"{ model | | Path to 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 (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) }");
|
||||
|
||||
static Mat postprocessOutput(const Mat &output, const Size &originalSize)
|
||||
{
|
||||
Mat squeezed;
|
||||
if (output.dims == 4 && output.size[0] == 1)
|
||||
{
|
||||
vector<int> newShape = {output.size[1], output.size[2], output.size[3]};
|
||||
squeezed = output.reshape(0, newShape);
|
||||
}
|
||||
else
|
||||
{
|
||||
squeezed = output.clone();
|
||||
}
|
||||
|
||||
Mat outputImage;
|
||||
vector<Mat> channels(3);
|
||||
for (int i = 0; i < 3; i++)
|
||||
{
|
||||
channels[2-i] = Mat(squeezed.size[1], squeezed.size[2], CV_32F,
|
||||
squeezed.ptr<float>(i));
|
||||
}
|
||||
merge(channels, outputImage);
|
||||
|
||||
outputImage = max(0.0, min(1.0, outputImage));
|
||||
outputImage.convertTo(outputImage, CV_8UC3, 255.0);
|
||||
|
||||
Size targetSize(originalSize.width * 4, originalSize.height * 4);
|
||||
Mat result;
|
||||
resize(outputImage, result, targetSize);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
static Mat applySuperResolution(Net &net, const Mat &image, float scale, const Scalar &mean, bool swapRB, int width, int height)
|
||||
{
|
||||
Mat blob = blobFromImage(image, scale, Size(width, height), mean, swapRB, false, CV_32F);
|
||||
|
||||
net.setInput(blob);
|
||||
Mat output;
|
||||
net.forward(output);
|
||||
|
||||
return postprocessOutput(output, Size(image.cols, image.rows));
|
||||
}
|
||||
|
||||
static double calculateFontScale(const Mat &image)
|
||||
{
|
||||
double baseScale = min(image.cols, image.rows) / 800.0;
|
||||
return max(0.5, baseScale);
|
||||
}
|
||||
|
||||
static void processFrame(Net &net, Mat &frame, float scale, const Scalar &mean, bool swapRB, int width, int height)
|
||||
{
|
||||
Mat result = applySuperResolution(net, frame, scale, mean, swapRB, width, height);
|
||||
|
||||
double fontScale = calculateFontScale(frame);
|
||||
int thickness = max(1, (int)(fontScale * 2));
|
||||
|
||||
putText(frame, "Original", Point(10, 30),
|
||||
FONT_HERSHEY_SIMPLEX, fontScale, Scalar(0, 255, 0), thickness);
|
||||
|
||||
double resultFontScale = calculateFontScale(result);
|
||||
int resultThickness = max(1, (int)(resultFontScale * 2));
|
||||
|
||||
putText(result, "Super-Resolution 4x", Point(20, 50),
|
||||
FONT_HERSHEY_SIMPLEX, resultFontScale, Scalar(0, 255, 0), resultThickness);
|
||||
|
||||
imshow("Input", frame);
|
||||
imshow("Super-Resolution", result);
|
||||
}
|
||||
|
||||
int main(int argc, char **argv)
|
||||
{
|
||||
const string about =
|
||||
"This sample demonstrates super-resolution using the SeeMore model.\n"
|
||||
"The model upscales images by 4x while enhancing details.\n\n"
|
||||
"Usage examples:\n"
|
||||
"\t./super_resolution\n"
|
||||
"\t./super_resolution --input=image.jpg\n"
|
||||
"\t./super_resolution --input=../data/chicky_512.png\n";
|
||||
|
||||
string keys = param_keys + backend_keys + target_keys;
|
||||
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
if (parser.has("help"))
|
||||
{
|
||||
cout << about << endl;
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
string modelName = parser.get<String>("@alias");
|
||||
string zooFile = samples::findFile(parser.get<String>("zoo"));
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile);
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
|
||||
float scale = parser.get<float>("scale");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
String backend = parser.get<String>("backend");
|
||||
String target = parser.get<String>("target");
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
string model = findModel(parser.get<String>("model"), sha1);
|
||||
int width = parser.get<int>("width");
|
||||
int height = parser.get<int>("height");
|
||||
string inputPath = findFile(parser.get<String>("input"));
|
||||
|
||||
if (model.empty())
|
||||
{
|
||||
cerr << "Model file not found" << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
Net net;
|
||||
try
|
||||
{
|
||||
net = readNetFromONNX(model);
|
||||
net.setPreferableBackend(getBackendID(backend));
|
||||
net.setPreferableTarget(getTargetID(target));
|
||||
}
|
||||
catch (const Exception &e)
|
||||
{
|
||||
cerr << "Error loading model: " << e.what() << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
Mat testImage = imread(inputPath);
|
||||
if (testImage.empty())
|
||||
{
|
||||
cerr << "Cannot load image: " << inputPath << endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
namedWindow("Input", WINDOW_NORMAL);
|
||||
namedWindow("Super-Resolution", WINDOW_NORMAL);
|
||||
moveWindow("Input", WINDOW_OFFSET_X, WINDOW_OFFSET_Y);
|
||||
moveWindow("Super-Resolution", WINDOW_OFFSET_X + testImage.cols + WINDOW_SPACING, WINDOW_OFFSET_Y);
|
||||
|
||||
processFrame(net, testImage, scale, mean, swapRB, width, height);
|
||||
waitKey(0);
|
||||
destroyAllWindows();
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
"""
|
||||
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.
|
||||
|
||||
|
||||
This sample demonstrates super-resolution using the SeeMoreDetails model.
|
||||
The model upscales images by 4x while enhancing details and reducing noise.
|
||||
Supports image inputs only.
|
||||
|
||||
SeeMoreDetails Repo: https://github.com/eduardzamfir/seemoredetails
|
||||
"""
|
||||
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import numpy as np
|
||||
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", help="Path to input image file.", default="chicky_512.png", 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()
|
||||
|
||||
model_name = "seemoredetails"
|
||||
add_preproc_args(args.zoo, parser, "super_resolution", model_name)
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
parents=[parser],
|
||||
description="""
|
||||
To run:
|
||||
Default image:
|
||||
python super_resolution.py
|
||||
Image processing:
|
||||
python super_resolution.py --input=path/to/your/input/image.jpg
|
||||
|
||||
The model performs 4x super-resolution on input images.
|
||||
""",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
return parser.parse_args(func_args)
|
||||
|
||||
def load_model(args):
|
||||
"""Load the super-resolution model"""
|
||||
try:
|
||||
model_path = findModel(args.model, args.sha1)
|
||||
net = cv.dnn.readNetFromONNX(model_path)
|
||||
net.setPreferableBackend(get_backend_id(args.backend))
|
||||
net.setPreferableTarget(get_target_id(args.target))
|
||||
return net
|
||||
except Exception as e:
|
||||
print(f"Error loading model: {e}")
|
||||
return None
|
||||
|
||||
def postprocess_output(output, args, original_shape=None):
|
||||
"""Postprocess model output to displayable image"""
|
||||
output = np.squeeze(output, axis=0)
|
||||
output = np.clip(output, 0, 1)
|
||||
output = np.transpose(output, (1, 2, 0))
|
||||
output = (output * 255).astype(np.uint8)
|
||||
|
||||
output = cv.cvtColor(output, cv.COLOR_RGB2BGR)
|
||||
|
||||
if original_shape is not None:
|
||||
target_height, target_width = original_shape
|
||||
upscaled_height, upscaled_width = target_height * 4, target_width * 4
|
||||
output = cv.resize(output, (upscaled_width, upscaled_height))
|
||||
|
||||
return output
|
||||
|
||||
def apply_super_resolution(net, image, args):
|
||||
"""Apply super-resolution to a single image"""
|
||||
original_shape = image.shape[:2]
|
||||
|
||||
blob = cv.dnn.blobFromImage(
|
||||
image,
|
||||
scalefactor=args.scale,
|
||||
size=(args.width, args.height),
|
||||
mean=args.mean,
|
||||
swapRB=args.rgb,
|
||||
crop=False,
|
||||
)
|
||||
|
||||
net.setInput(blob)
|
||||
t0 = cv.getTickCount()
|
||||
output = net.forward()
|
||||
t = (cv.getTickCount() - t0) / cv.getTickFrequency()
|
||||
|
||||
result = postprocess_output(output, args, original_shape)
|
||||
|
||||
label = "Inference time: %.2f ms" % (t * 1000.0)
|
||||
cv.putText(result, label, (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
|
||||
return result
|
||||
|
||||
def main(func_args=None):
|
||||
args = get_args_parser(func_args)
|
||||
|
||||
net = load_model(args)
|
||||
if net is None:
|
||||
print("Failed to load model.")
|
||||
return -1
|
||||
|
||||
input_path = cv.samples.findFile(args.input)
|
||||
image = cv.imread(input_path)
|
||||
if image is None:
|
||||
print(f"Cannot load image: {input_path}")
|
||||
return -1
|
||||
|
||||
print(f"Processing image: {input_path}")
|
||||
result = apply_super_resolution(net, image, args)
|
||||
|
||||
cv.namedWindow("Input", cv.WINDOW_NORMAL)
|
||||
cv.namedWindow("Super-Resolution Result", cv.WINDOW_NORMAL)
|
||||
cv.imshow("Input", image)
|
||||
cv.imshow("Super-Resolution Result", result)
|
||||
print("Press 'q' to quit...")
|
||||
while True:
|
||||
key = cv.waitKey(0) & 0xFF
|
||||
if key == ord("q"):
|
||||
break
|
||||
cv.destroyAllWindows()
|
||||
return 0
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,262 @@
|
||||
/*
|
||||
Text detection model (EAST): https://github.com/argman/EAST
|
||||
Download link for EAST model: https://www.dropbox.com/s/r2ingd0l3zt8hxs/frozen_east_text_detection.tar.gz?dl=1
|
||||
|
||||
DB detector model:
|
||||
https://drive.google.com/uc?export=download&id=17_ABp79PlFt9yPCxSaarVc_DKTmrSGGf
|
||||
|
||||
CRNN Text recognition model sourced from: https://github.com/meijieru/crnn.pytorch
|
||||
How to convert from .pb to .onnx:
|
||||
Using classes from: https://github.com/meijieru/crnn.pytorch/blob/master/models/crnn.py
|
||||
|
||||
Additional converted ONNX text recognition models available for direct download:
|
||||
Download link: https://drive.google.com/drive/folders/1cTbQ3nuZG-EKWak6emD_s8_hHXWz7lAr?usp=sharing
|
||||
These models are taken from: https://github.com/clovaai/deep-text-recognition-benchmark
|
||||
|
||||
Importing and using the CRNN model in PyTorch:
|
||||
import torch
|
||||
from models.crnn import CRNN
|
||||
|
||||
model = CRNN(32, 1, 37, 256)
|
||||
model.load_state_dict(torch.load('crnn.pth'))
|
||||
dummy_input = torch.randn(1, 1, 32, 100)
|
||||
torch.onnx.export(model, dummy_input, "crnn.onnx", verbose=True)
|
||||
|
||||
Usage: ./example_dnn_text_detection DB
|
||||
*/
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
|
||||
#include <opencv2/geometry.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/highgui.hpp>
|
||||
#include <opencv2/dnn.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
using namespace cv;
|
||||
using namespace std;
|
||||
using namespace cv::dnn;
|
||||
|
||||
const string about = "Use this script for Text Detection and Recognition using OpenCV. \n\n"
|
||||
"Firstly, download required models 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_text_detection modelName(i.e. DB or East) --ocr_model=<path to VGG_CTC.onnx>\n\n"
|
||||
"Detection model path can also be specified using --model argument. \n\n"
|
||||
"Download ocr model using: python download_models.py OCR \n\n";
|
||||
|
||||
// Command-line keys to parse the input arguments
|
||||
string keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ input i | right.jpg | Path to an input image. }"
|
||||
"{ @alias | | 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 }"
|
||||
"{ ocr_model | | Path to a binary .onnx model for recognition. }"
|
||||
"{ model | | Path to detection model file. }"
|
||||
"{ thr | 0.5 | Confidence threshold for EAST detector. }"
|
||||
"{ nms | 0.4 | Non-maximum suppression threshold for EAST detector. }"
|
||||
"{ binaryThreshold bt | 0.3 | Confidence threshold for the binary map in DB detector. }"
|
||||
"{ polygonThreshold pt | 0.5 | Confidence threshold for polygons in DB detector. }"
|
||||
"{ maxCandidate max | 200 | Max candidates for polygons in DB detector. }"
|
||||
"{ unclipRatio ratio | 2.0 | Unclip ratio for DB detector. }"
|
||||
"{ vocabularyPath vp | alphabet_36.txt | Path to vocabulary file. }";
|
||||
|
||||
// Function prototype for the four-point perspective transform
|
||||
static void fourPointsTransform(const Mat& frame, const Point2f vertices[], Mat& result);
|
||||
static void processFrame(
|
||||
const Mat& frame,
|
||||
const vector<vector<Point>>& detResults,
|
||||
const std::string& ocr_model,
|
||||
bool imreadRGB,
|
||||
Mat& board,
|
||||
FontFace& fontFace,
|
||||
int fontSize,
|
||||
int fontWeight,
|
||||
const vector<std::string>& vocabulary
|
||||
);
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
// Setting up command-line parser with the specified keys
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
|
||||
if (!parser.has("@alias") || parser.has("help"))
|
||||
{
|
||||
cout << about << endl;
|
||||
parser.printMessage();
|
||||
return -1;
|
||||
}
|
||||
const string modelName = parser.get<String>("@alias");
|
||||
const string zooFile = findFile(parser.get<String>("zoo"));
|
||||
|
||||
keys += genPreprocArguments(modelName, zooFile, "");
|
||||
keys += genPreprocArguments(modelName, zooFile, "ocr_");
|
||||
parser = CommandLineParser(argc, argv, keys);
|
||||
parser.about(about);
|
||||
|
||||
// Parsing command-line arguments
|
||||
|
||||
String sha1 = parser.get<String>("sha1");
|
||||
String ocr_sha1 = parser.get<String>("ocr_sha1");
|
||||
String detModelPath = findModel(parser.get<String>("model"), sha1);
|
||||
String ocr = findModel(parser.get<String>("ocr_model"), ocr_sha1);
|
||||
int height = parser.get<int>("height");
|
||||
int width = parser.get<int>("width");
|
||||
bool imreadRGB = parser.get<bool>("rgb");
|
||||
String vocPath = parser.get<String>("vocabularyPath");
|
||||
float binThresh = parser.get<float>("binaryThreshold");
|
||||
float polyThresh = parser.get<float>("polygonThreshold");
|
||||
double unclipRatio = parser.get<double>("unclipRatio");
|
||||
uint maxCandidates = parser.get<uint>("maxCandidate");
|
||||
float confThreshold = parser.get<float>("thr");
|
||||
float nmsThreshold = parser.get<float>("nms");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
|
||||
// Ensuring the provided arguments are valid
|
||||
if (!parser.check()) {
|
||||
parser.printErrors();
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Asserting detection model path is provided
|
||||
CV_Assert(!detModelPath.empty());
|
||||
|
||||
vector<vector<Point>> detResults;
|
||||
// Reading the input image
|
||||
Mat frame = imread(samples::findFile(parser.get<String>("input")));
|
||||
Mat board(frame.size(), frame.type(), Scalar(255, 255, 255));
|
||||
int stdSize = 20;
|
||||
int stdWeight = 400;
|
||||
int stdImgSize = 512;
|
||||
int imgWidth = min(frame.rows, frame.cols);
|
||||
int size = (stdSize*imgWidth)/stdImgSize;
|
||||
int weight = (stdWeight*imgWidth)/stdImgSize;
|
||||
FontFace fontFace("sans");
|
||||
|
||||
// Initializing and configuring the text detection model based on the provided config
|
||||
if (modelName == "East") {
|
||||
// EAST Detector initialization
|
||||
TextDetectionModel_EAST detector(detModelPath);
|
||||
detector.setConfidenceThreshold(confThreshold)
|
||||
.setNMSThreshold(nmsThreshold);
|
||||
// Setting input parameters specific to EAST model
|
||||
detector.setInputParams(1.0, Size(width, height), mean, true);
|
||||
// Performing text detection
|
||||
detector.detect(frame, detResults);
|
||||
}
|
||||
else if (modelName == "DB") {
|
||||
// DB Detector initialization
|
||||
TextDetectionModel_DB detector(detModelPath);
|
||||
detector.setBinaryThreshold(binThresh)
|
||||
.setPolygonThreshold(polyThresh)
|
||||
.setUnclipRatio(unclipRatio)
|
||||
.setMaxCandidates(maxCandidates);
|
||||
// Setting input parameters specific to DB model
|
||||
detector.setInputParams(1.0 / 255.0, Size(width, height), mean);
|
||||
// Performing text detection
|
||||
detector.detect(frame, detResults);
|
||||
}
|
||||
else {
|
||||
cout << "[ERROR]: Unsupported file config for the detector model. Valid values: east/db" << endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Reading and storing vocabulary for text recognition
|
||||
CV_Assert(!vocPath.empty());
|
||||
ifstream vocFile;
|
||||
vocFile.open(samples::findFile(vocPath));
|
||||
CV_Assert(vocFile.is_open());
|
||||
std::string vocLine;
|
||||
vector<std::string> vocabulary;
|
||||
while (getline(vocFile, vocLine)) {
|
||||
vocabulary.push_back(vocLine);
|
||||
}
|
||||
|
||||
processFrame(frame, detResults, ocr, imreadRGB, board, fontFace, size, weight, vocabulary);
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Performs a perspective transform for a four-point region
|
||||
static void fourPointsTransform(const Mat& frame, const Point2f vertices[], Mat& result) {
|
||||
const Size outputSize = Size(100, 32);
|
||||
// Defining target vertices for the perspective transform
|
||||
Point2f targetVertices[4] = {
|
||||
Point(0, outputSize.height - 1),
|
||||
Point(0, 0),
|
||||
Point(outputSize.width - 1, 0),
|
||||
Point(outputSize.width - 1, outputSize.height - 1)
|
||||
};
|
||||
// Computing the perspective transform matrix
|
||||
Mat rotationMatrix = getPerspectiveTransform(vertices, targetVertices);
|
||||
// Applying the perspective transform to the region
|
||||
warpPerspective(frame, result, rotationMatrix, outputSize);
|
||||
}
|
||||
|
||||
void processFrame(
|
||||
const Mat& frame,
|
||||
const vector<vector<Point>>& detResults,
|
||||
const std::string& ocr_model,
|
||||
bool imreadRGB,
|
||||
Mat& board,
|
||||
FontFace& fontFace,
|
||||
int fontSize,
|
||||
int fontWeight,
|
||||
const vector<std::string>& vocabulary
|
||||
) {
|
||||
if (detResults.size() > 0) {
|
||||
// Text Recognition
|
||||
Mat recInput;
|
||||
if (!imreadRGB) {
|
||||
cvtColor(frame, recInput, cv::COLOR_BGR2GRAY);
|
||||
} else {
|
||||
recInput = frame;
|
||||
}
|
||||
|
||||
vector<vector<Point>> contours;
|
||||
for (uint i = 0; i < detResults.size(); i++) {
|
||||
const auto& quadrangle = detResults[i];
|
||||
CV_CheckEQ(quadrangle.size(), (size_t)4, "");
|
||||
|
||||
contours.emplace_back(quadrangle);
|
||||
|
||||
vector<Point2f> quadrangle_2f;
|
||||
for (int j = 0; j < 4; j++)
|
||||
quadrangle_2f.emplace_back(detResults[i][j]);
|
||||
|
||||
// Cropping the detected text region using a four-point transform
|
||||
Mat cropped;
|
||||
fourPointsTransform(recInput, &quadrangle_2f[0], cropped);
|
||||
|
||||
if(!ocr_model.empty()){
|
||||
TextRecognitionModel recognizer(ocr_model);
|
||||
recognizer.setVocabulary(vocabulary);
|
||||
recognizer.setDecodeType("CTC-greedy");
|
||||
|
||||
// Setting input parameters for the recognition model
|
||||
double recScale = 1.0 / 127.5;
|
||||
Scalar recMean = Scalar(127.5);
|
||||
Size recInputSize = Size(100, 32);
|
||||
recognizer.setInputParams(recScale, recInputSize, recMean);
|
||||
// Recognizing text from the cropped image
|
||||
string recognitionResult = recognizer.recognize(cropped);
|
||||
cout << i << ": '" << recognitionResult << "'" << endl;
|
||||
|
||||
// Displaying the recognized text on the image
|
||||
putText(board, recognitionResult, Point(detResults[i][1].x, detResults[i][0].y), Scalar(0, 0, 0), fontFace, fontSize, fontWeight);
|
||||
}
|
||||
else{
|
||||
cout << "[WARN] Please pass the path to the ocr model using --ocr_model to get the recognised text." << endl;
|
||||
}
|
||||
}
|
||||
// Drawing detected text regions on the image
|
||||
polylines(board, contours, true, Scalar(200, 255, 200), 1);
|
||||
polylines(frame, contours, true, Scalar(0, 255, 0), 1);
|
||||
} else {
|
||||
cout << "No Text Detected." << endl;
|
||||
}
|
||||
|
||||
// Displaying the final image with detected and recognized text
|
||||
Mat stacked;
|
||||
hconcat(frame, board, stacked);
|
||||
imshow("Text Detection and Recognition", stacked);
|
||||
waitKey(0);
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
'''
|
||||
Text detection model (EAST): https://github.com/argman/EAST
|
||||
Download link for EAST model: https://www.dropbox.com/s/r2ingd0l3zt8hxs/frozen_east_text_detection.tar.gz?dl=1
|
||||
|
||||
DB detector model:
|
||||
https://drive.google.com/uc?export=download&id=17_ABp79PlFt9yPCxSaarVc_DKTmrSGGf
|
||||
|
||||
CRNN Text recognition model sourced from: https://github.com/meijieru/crnn.pytorch
|
||||
How to convert from .pb to .onnx:
|
||||
Using classes from: https://github.com/meijieru/crnn.pytorch/blob/master/models/crnn.py
|
||||
|
||||
Additional converted ONNX text recognition models available for direct download:
|
||||
Download link: https://drive.google.com/drive/folders/1cTbQ3nuZG-EKWak6emD_s8_hHXWz7lAr?usp=sharing
|
||||
These models are taken from: https://github.com/clovaai/deep-text-recognition-benchmark
|
||||
|
||||
Importing and using the CRNN model in PyTorch:
|
||||
import torch
|
||||
from models.crnn import CRNN
|
||||
|
||||
model = CRNN(32, 1, 37, 256)
|
||||
model.load_state_dict(torch.load('crnn.pth'))
|
||||
dummy_input = torch.randn(1, 1, 32, 100)
|
||||
torch.onnx.export(model, dummy_input, "crnn.onnx", verbose=True)
|
||||
|
||||
Usage: python text_detection.py DB --ocr_model=<path to recognition model>
|
||||
|
||||
'''
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
from common import *
|
||||
|
||||
def help():
|
||||
print(
|
||||
'''
|
||||
Use this script for Text Detection and Recognition using OpenCV.
|
||||
|
||||
Firstly, download required models 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.
|
||||
|
||||
Example: python download_models.py East
|
||||
python download_models.py OCR
|
||||
|
||||
To run:
|
||||
Example: python text_detection.py modelName(i.e. DB or East)
|
||||
|
||||
Detection model path can also be specified using --model argument and ocr model can be specified using --ocr_model.
|
||||
'''
|
||||
)
|
||||
|
||||
############ Add argument parser for command line arguments ############
|
||||
def get_args_parser():
|
||||
parser = argparse.ArgumentParser(add_help=False)
|
||||
parser.add_argument('--input', default='right.jpg',
|
||||
help='Path to input image or video file. Skip this argument to capture frames from a camera.')
|
||||
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('--thr', type=float, default=0.5,
|
||||
help='Confidence threshold.')
|
||||
parser.add_argument('--nms', type=float, default=0.4,
|
||||
help='Non-maximum suppression threshold.')
|
||||
parser.add_argument('--binary_threshold', type=float, default=0.3,
|
||||
help='Confidence threshold for the binary map in DB detector. ')
|
||||
parser.add_argument('--polygon_threshold', type=float, default=0.5,
|
||||
help='Confidence threshold for polygons in DB detector.')
|
||||
parser.add_argument('--max_candidate', type=int, default=200,
|
||||
help='Max candidates for polygons in DB detector.')
|
||||
parser.add_argument('--unclip_ratio', type=float, default=2.0,
|
||||
help='Unclip ratio for DB detector.')
|
||||
parser.add_argument('--vocabulary_path', default='alphabet_36.txt',
|
||||
help='Path to vocabulary file.')
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
add_preproc_args(args.zoo, parser, 'text_detection', prefix="")
|
||||
add_preproc_args(args.zoo, parser, 'text_recognition', prefix="ocr_")
|
||||
parser = argparse.ArgumentParser(parents=[parser],
|
||||
description='Text Detection and Recognition using OpenCV.',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
return parser.parse_args()
|
||||
|
||||
def fourPointsTransform(frame, vertices):
|
||||
vertices = np.asarray(vertices)
|
||||
outputSize = (100, 32)
|
||||
targetVertices = np.array([
|
||||
[0, outputSize[1] - 1],
|
||||
[0, 0],
|
||||
[outputSize[0] - 1, 0],
|
||||
[outputSize[0] - 1, outputSize[1] - 1]], dtype="float32")
|
||||
|
||||
rotationMatrix = cv2.getPerspectiveTransform(vertices, targetVertices)
|
||||
result = cv2.warpPerspective(frame, rotationMatrix, outputSize)
|
||||
return result
|
||||
|
||||
def main():
|
||||
args = get_args_parser()
|
||||
if args.alias is None or hasattr(args, 'help'):
|
||||
help()
|
||||
exit(1)
|
||||
|
||||
args.model = findModel(args.model, args.sha1)
|
||||
|
||||
args.ocr_model = findModel(args.ocr_model, args.ocr_sha1)
|
||||
args.input = findFile(args.input)
|
||||
args.vocabulary_path = findFile(args.vocabulary_path)
|
||||
|
||||
frame = cv2.imread(args.input)
|
||||
board = np.ones_like(frame)*255
|
||||
|
||||
stdSize = 0.8
|
||||
stdWeight = 2
|
||||
stdImgSize = 512
|
||||
imgWidth = min(frame.shape[:2])
|
||||
fontSize = (stdSize*imgWidth)/stdImgSize
|
||||
fontThickness = max(1,(stdWeight*imgWidth)//stdImgSize)
|
||||
|
||||
if(args.alias == "DB"):
|
||||
# DB Detector initialization
|
||||
detector = cv2.dnn_TextDetectionModel_DB(args.model)
|
||||
detector.setBinaryThreshold(args.binary_threshold)
|
||||
detector.setPolygonThreshold(args.polygon_threshold)
|
||||
detector.setUnclipRatio(args.unclip_ratio)
|
||||
detector.setMaxCandidates(args.max_candidate)
|
||||
# Setting input parameters specific to the DB model
|
||||
detector.setInputParams(scale=args.scale, size=(args.width, args.height), mean=args.mean)
|
||||
# Performing text detection
|
||||
detResults = detector.detect(frame)
|
||||
elif(args.alias == "East"):
|
||||
# EAST Detector initialization
|
||||
detector = cv2.dnn_TextDetectionModel_EAST(args.model)
|
||||
detector.setConfidenceThreshold(args.thr)
|
||||
detector.setNMSThreshold(args.nms)
|
||||
# Setting input parameters specific to EAST model
|
||||
detector.setInputParams(scale=args.scale, size=(args.width, args.height), mean=args.mean, swapRB=True)
|
||||
# Perfroming text detection
|
||||
detResults = detector.detect(frame)
|
||||
|
||||
# Open the vocabulary file and read lines into a list
|
||||
with open(args.vocabulary_path, 'r') as voc_file:
|
||||
vocabulary = [line.strip() for line in voc_file]
|
||||
|
||||
if args.ocr_model is None:
|
||||
print("[ERROR] Please pass the path to the ocr model using --ocr_model to run the sample")
|
||||
exit(1)
|
||||
# Initialize the text recognition model with the specified model path
|
||||
recognizer = cv2.dnn_TextRecognitionModel(args.ocr_model)
|
||||
|
||||
# Set the vocabulary for the model
|
||||
recognizer.setVocabulary(vocabulary)
|
||||
|
||||
# Set the decoding method to 'CTC-greedy'
|
||||
recognizer.setDecodeType("CTC-greedy")
|
||||
|
||||
recScale = 1.0 / 127.5
|
||||
recMean = (127.5, 127.5, 127.5)
|
||||
recInputSize = (100, 32)
|
||||
recognizer.setInputParams(scale=recScale, size=recInputSize, mean=recMean)
|
||||
|
||||
if len(detResults) > 0:
|
||||
recInput = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if not args.rgb else frame.copy()
|
||||
contours = []
|
||||
|
||||
for i, (quadrangle, _) in enumerate(zip(detResults[0], detResults[1])):
|
||||
if isinstance(quadrangle, np.ndarray):
|
||||
quadrangle = np.array(quadrangle).astype(np.float32)
|
||||
|
||||
if quadrangle is None or len(quadrangle) != 4:
|
||||
print("Skipping a quadrangle with incorrect points or transformation failed.")
|
||||
continue
|
||||
|
||||
contours.append(np.array(quadrangle, dtype=np.int32))
|
||||
cropped = fourPointsTransform(recInput, quadrangle)
|
||||
recognitionResult = recognizer.recognize(cropped)
|
||||
print(f"{i}: '{recognitionResult}'")
|
||||
|
||||
try:
|
||||
text_origin = (int(quadrangle[1][0]), int(quadrangle[0][1]))
|
||||
cv2.putText(board, recognitionResult, text_origin, cv2.FONT_HERSHEY_SIMPLEX, fontSize, (0, 0, 0), fontThickness)
|
||||
except Exception as e:
|
||||
print("Failed to write text on the frame:", e)
|
||||
else:
|
||||
print("Skipping a detection with invalid format:", quadrangle)
|
||||
|
||||
cv2.polylines(frame, contours, True, (0, 255, 0), 1)
|
||||
cv2.polylines(board, contours, True, (200, 255, 200), 1)
|
||||
else:
|
||||
print("No Text Detected.")
|
||||
|
||||
stacked = cv2.hconcat([frame, board])
|
||||
cv2.imshow("Text Detection and Recognition", stacked)
|
||||
cv2.waitKey(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,333 @@
|
||||
def tokenize(s):
|
||||
tokens = []
|
||||
token = ""
|
||||
isString = False
|
||||
isComment = False
|
||||
for symbol in s:
|
||||
isComment = (isComment and symbol != '\n') or (not isString and symbol == '#')
|
||||
if isComment:
|
||||
continue
|
||||
|
||||
if symbol == ' ' or symbol == '\t' or symbol == '\r' or symbol == '\'' or \
|
||||
symbol == '\n' or symbol == ':' or symbol == '\"' or symbol == ';' or \
|
||||
symbol == ',':
|
||||
|
||||
if (symbol == '\"' or symbol == '\'') and isString:
|
||||
tokens.append(token)
|
||||
token = ""
|
||||
else:
|
||||
if isString:
|
||||
token += symbol
|
||||
elif token:
|
||||
tokens.append(token)
|
||||
token = ""
|
||||
isString = (symbol == '\"' or symbol == '\'') ^ isString
|
||||
|
||||
elif symbol == '{' or symbol == '}' or symbol == '[' or symbol == ']':
|
||||
if token:
|
||||
tokens.append(token)
|
||||
token = ""
|
||||
tokens.append(symbol)
|
||||
else:
|
||||
token += symbol
|
||||
if token:
|
||||
tokens.append(token)
|
||||
return tokens
|
||||
|
||||
|
||||
def parseMessage(tokens, idx):
|
||||
msg = {}
|
||||
assert(tokens[idx] == '{')
|
||||
|
||||
isArray = False
|
||||
while True:
|
||||
if not isArray:
|
||||
idx += 1
|
||||
if idx < len(tokens):
|
||||
fieldName = tokens[idx]
|
||||
else:
|
||||
return None
|
||||
if fieldName == '}':
|
||||
break
|
||||
|
||||
idx += 1
|
||||
fieldValue = tokens[idx]
|
||||
|
||||
if fieldValue == '{':
|
||||
embeddedMsg, idx = parseMessage(tokens, idx)
|
||||
if fieldName in msg:
|
||||
msg[fieldName].append(embeddedMsg)
|
||||
else:
|
||||
msg[fieldName] = [embeddedMsg]
|
||||
elif fieldValue == '[':
|
||||
isArray = True
|
||||
elif fieldValue == ']':
|
||||
isArray = False
|
||||
else:
|
||||
if fieldName in msg:
|
||||
msg[fieldName].append(fieldValue)
|
||||
else:
|
||||
msg[fieldName] = [fieldValue]
|
||||
return msg, idx
|
||||
|
||||
|
||||
def readTextMessage(filePath):
|
||||
if not filePath:
|
||||
return {}
|
||||
with open(filePath, 'rt') as f:
|
||||
content = f.read()
|
||||
|
||||
tokens = tokenize('{' + content + '}')
|
||||
msg = parseMessage(tokens, 0)
|
||||
return msg[0] if msg else {}
|
||||
|
||||
|
||||
def listToTensor(values):
|
||||
if all([isinstance(v, float) for v in values]):
|
||||
dtype = 'DT_FLOAT'
|
||||
field = 'float_val'
|
||||
elif all([isinstance(v, int) for v in values]):
|
||||
dtype = 'DT_INT32'
|
||||
field = 'int_val'
|
||||
else:
|
||||
raise Exception('Wrong values types')
|
||||
|
||||
msg = {
|
||||
'tensor': {
|
||||
'dtype': dtype,
|
||||
'tensor_shape': {
|
||||
'dim': {
|
||||
'size': len(values)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
msg['tensor'][field] = values
|
||||
return msg
|
||||
|
||||
|
||||
def addConstNode(name, values, graph_def):
|
||||
node = NodeDef()
|
||||
node.name = name
|
||||
node.op = 'Const'
|
||||
node.addAttr('value', values)
|
||||
graph_def.node.extend([node])
|
||||
|
||||
|
||||
def addSlice(inp, out, begins, sizes, graph_def):
|
||||
beginsNode = NodeDef()
|
||||
beginsNode.name = out + '/begins'
|
||||
beginsNode.op = 'Const'
|
||||
beginsNode.addAttr('value', begins)
|
||||
graph_def.node.extend([beginsNode])
|
||||
|
||||
sizesNode = NodeDef()
|
||||
sizesNode.name = out + '/sizes'
|
||||
sizesNode.op = 'Const'
|
||||
sizesNode.addAttr('value', sizes)
|
||||
graph_def.node.extend([sizesNode])
|
||||
|
||||
sliced = NodeDef()
|
||||
sliced.name = out
|
||||
sliced.op = 'Slice'
|
||||
sliced.input.append(inp)
|
||||
sliced.input.append(beginsNode.name)
|
||||
sliced.input.append(sizesNode.name)
|
||||
graph_def.node.extend([sliced])
|
||||
|
||||
|
||||
def addReshape(inp, out, shape, graph_def):
|
||||
shapeNode = NodeDef()
|
||||
shapeNode.name = out + '/shape'
|
||||
shapeNode.op = 'Const'
|
||||
shapeNode.addAttr('value', shape)
|
||||
graph_def.node.extend([shapeNode])
|
||||
|
||||
reshape = NodeDef()
|
||||
reshape.name = out
|
||||
reshape.op = 'Reshape'
|
||||
reshape.input.append(inp)
|
||||
reshape.input.append(shapeNode.name)
|
||||
graph_def.node.extend([reshape])
|
||||
|
||||
|
||||
def addSoftMax(inp, out, graph_def):
|
||||
softmax = NodeDef()
|
||||
softmax.name = out
|
||||
softmax.op = 'Softmax'
|
||||
softmax.addAttr('axis', -1)
|
||||
softmax.input.append(inp)
|
||||
graph_def.node.extend([softmax])
|
||||
|
||||
|
||||
def addFlatten(inp, out, graph_def):
|
||||
flatten = NodeDef()
|
||||
flatten.name = out
|
||||
flatten.op = 'Flatten'
|
||||
flatten.input.append(inp)
|
||||
graph_def.node.extend([flatten])
|
||||
|
||||
|
||||
class NodeDef:
|
||||
def __init__(self):
|
||||
self.input = []
|
||||
self.name = ""
|
||||
self.op = ""
|
||||
self.attr = {}
|
||||
|
||||
def addAttr(self, key, value):
|
||||
assert(not key in self.attr)
|
||||
if isinstance(value, bool):
|
||||
self.attr[key] = {'b': value}
|
||||
elif isinstance(value, int):
|
||||
self.attr[key] = {'i': value}
|
||||
elif isinstance(value, float):
|
||||
self.attr[key] = {'f': value}
|
||||
elif isinstance(value, str):
|
||||
self.attr[key] = {'s': value}
|
||||
elif isinstance(value, list):
|
||||
self.attr[key] = listToTensor(value)
|
||||
else:
|
||||
raise Exception('Unknown type of attribute ' + key)
|
||||
|
||||
def Clear(self):
|
||||
self.input = []
|
||||
self.name = ""
|
||||
self.op = ""
|
||||
self.attr = {}
|
||||
|
||||
|
||||
class GraphDef:
|
||||
def __init__(self):
|
||||
self.node = []
|
||||
|
||||
def save(self, filePath):
|
||||
with open(filePath, 'wt') as f:
|
||||
|
||||
def printAttr(d, indent):
|
||||
indent = ' ' * indent
|
||||
for key, value in sorted(d.items(), key=lambda x:x[0].lower()):
|
||||
value = value if isinstance(value, list) else [value]
|
||||
for v in value:
|
||||
if isinstance(v, dict):
|
||||
f.write(indent + key + ' {\n')
|
||||
printAttr(v, len(indent) + 2)
|
||||
f.write(indent + '}\n')
|
||||
else:
|
||||
isString = False
|
||||
if isinstance(v, str) and not v.startswith('DT_'):
|
||||
try:
|
||||
float(v)
|
||||
except:
|
||||
isString = True
|
||||
|
||||
if isinstance(v, bool):
|
||||
printed = 'true' if v else 'false'
|
||||
elif v == 'true' or v == 'false':
|
||||
printed = 'true' if v == 'true' else 'false'
|
||||
elif isString:
|
||||
printed = '\"%s\"' % v
|
||||
else:
|
||||
printed = str(v)
|
||||
f.write(indent + key + ': ' + printed + '\n')
|
||||
|
||||
for node in self.node:
|
||||
f.write('node {\n')
|
||||
f.write(' name: \"%s\"\n' % node.name)
|
||||
f.write(' op: \"%s\"\n' % node.op)
|
||||
for inp in node.input:
|
||||
f.write(' input: \"%s\"\n' % inp)
|
||||
for key, value in sorted(node.attr.items(), key=lambda x:x[0].lower()):
|
||||
f.write(' attr {\n')
|
||||
f.write(' key: \"%s\"\n' % key)
|
||||
f.write(' value {\n')
|
||||
printAttr(value, 6)
|
||||
f.write(' }\n')
|
||||
f.write(' }\n')
|
||||
f.write('}\n')
|
||||
|
||||
|
||||
def parseTextGraph(filePath):
|
||||
msg = readTextMessage(filePath)
|
||||
|
||||
graph = GraphDef()
|
||||
for node in msg['node']:
|
||||
graphNode = NodeDef()
|
||||
graphNode.name = node['name'][0]
|
||||
graphNode.op = node['op'][0]
|
||||
graphNode.input = node['input'] if 'input' in node else []
|
||||
|
||||
if 'attr' in node:
|
||||
for attr in node['attr']:
|
||||
graphNode.attr[attr['key'][0]] = attr['value'][0]
|
||||
|
||||
graph.node.append(graphNode)
|
||||
return graph
|
||||
|
||||
|
||||
# Removes Identity nodes
|
||||
def removeIdentity(graph_def):
|
||||
identities = {}
|
||||
for node in graph_def.node:
|
||||
if node.op == 'Identity' or node.op == 'IdentityN':
|
||||
inp = node.input[0]
|
||||
if inp in identities:
|
||||
identities[node.name] = identities[inp]
|
||||
else:
|
||||
identities[node.name] = inp
|
||||
graph_def.node.remove(node)
|
||||
|
||||
for node in graph_def.node:
|
||||
for i in range(len(node.input)):
|
||||
if node.input[i] in identities:
|
||||
node.input[i] = identities[node.input[i]]
|
||||
|
||||
|
||||
def removeUnusedNodesAndAttrs(to_remove, graph_def):
|
||||
unusedAttrs = ['T', 'Tshape', 'N', 'Tidx', 'Tdim', 'use_cudnn_on_gpu',
|
||||
'Index', 'Tperm', 'is_training', 'Tpaddings']
|
||||
|
||||
removedNodes = []
|
||||
|
||||
for i in reversed(range(len(graph_def.node))):
|
||||
op = graph_def.node[i].op
|
||||
name = graph_def.node[i].name
|
||||
|
||||
if to_remove(name, op):
|
||||
if op != 'Const':
|
||||
removedNodes.append(name)
|
||||
|
||||
del graph_def.node[i]
|
||||
else:
|
||||
for attr in unusedAttrs:
|
||||
if attr in graph_def.node[i].attr:
|
||||
del graph_def.node[i].attr[attr]
|
||||
|
||||
# Remove references to removed nodes except Const nodes.
|
||||
for node in graph_def.node:
|
||||
for i in reversed(range(len(node.input))):
|
||||
if node.input[i] in removedNodes:
|
||||
del node.input[i]
|
||||
|
||||
|
||||
def writeTextGraph(modelPath, outputPath, outNodes):
|
||||
try:
|
||||
import cv2 as cv
|
||||
|
||||
cv.dnn.writeTextGraph(modelPath, outputPath)
|
||||
except:
|
||||
import tensorflow as tf
|
||||
from tensorflow.tools.graph_transforms import TransformGraph
|
||||
|
||||
with tf.gfile.FastGFile(modelPath, 'rb') as f:
|
||||
graph_def = tf.GraphDef()
|
||||
graph_def.ParseFromString(f.read())
|
||||
|
||||
graph_def = TransformGraph(graph_def, ['image_tensor'], outNodes, ['sort_by_execution_order'])
|
||||
|
||||
for node in graph_def.node:
|
||||
if node.op == 'Const':
|
||||
if 'value' in node.attr and node.attr['value'].tensor.tensor_content:
|
||||
node.attr['value'].tensor.tensor_content = b''
|
||||
|
||||
tf.train.write_graph(graph_def, "", outputPath, as_text=True)
|
||||
@@ -0,0 +1,236 @@
|
||||
# This file is a part of OpenCV project.
|
||||
# It is a 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) 2020, Intel Corporation, all rights reserved.
|
||||
# Third party copyrights are property of their respective owners.
|
||||
#
|
||||
# Use this script to get the text graph representation (.pbtxt) of EfficientDet
|
||||
# deep learning network trained in https://github.com/google/automl.
|
||||
# Then you can import it with a binary frozen graph (.pb) using readNetFromTensorflow() function.
|
||||
# See details and examples on the following wiki page: https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API
|
||||
import argparse
|
||||
import re
|
||||
from math import sqrt
|
||||
from tf_text_graph_common import *
|
||||
|
||||
|
||||
class AnchorGenerator:
|
||||
def __init__(self, min_level, aspect_ratios, num_scales, anchor_scale):
|
||||
self.min_level = min_level
|
||||
self.aspect_ratios = aspect_ratios
|
||||
self.anchor_scale = anchor_scale
|
||||
self.scales = [2**(float(s) / num_scales) for s in range(num_scales)]
|
||||
|
||||
def get(self, layer_id):
|
||||
widths = []
|
||||
heights = []
|
||||
for s in self.scales:
|
||||
for a in self.aspect_ratios:
|
||||
base_anchor_size = 2**(self.min_level + layer_id) * self.anchor_scale
|
||||
heights.append(base_anchor_size * s * a[1])
|
||||
widths.append(base_anchor_size * s * a[0])
|
||||
return widths, heights
|
||||
|
||||
|
||||
def createGraph(modelPath, outputPath, min_level, aspect_ratios, num_scales,
|
||||
anchor_scale, num_classes, image_width, image_height):
|
||||
print('Min level: %d' % min_level)
|
||||
print('Anchor scale: %f' % anchor_scale)
|
||||
print('Num scales: %d' % num_scales)
|
||||
print('Aspect ratios: %s' % str(aspect_ratios))
|
||||
print('Number of classes: %d' % num_classes)
|
||||
print('Input image size: %dx%d' % (image_width, image_height))
|
||||
|
||||
# Read the graph.
|
||||
_inpNames = ['image_arrays']
|
||||
outNames = ['detections']
|
||||
|
||||
writeTextGraph(modelPath, outputPath, outNames)
|
||||
graph_def = parseTextGraph(outputPath)
|
||||
|
||||
def getUnconnectedNodes():
|
||||
unconnected = []
|
||||
for node in graph_def.node:
|
||||
if node.op == 'Const':
|
||||
continue
|
||||
unconnected.append(node.name)
|
||||
for inp in node.input:
|
||||
if inp in unconnected:
|
||||
unconnected.remove(inp)
|
||||
return unconnected
|
||||
|
||||
|
||||
nodesToKeep = ['truediv'] # Keep preprocessing nodes
|
||||
|
||||
removeIdentity(graph_def)
|
||||
|
||||
scopesToKeep = ('image_arrays', 'efficientnet', 'resample_p6', 'resample_p7',
|
||||
'fpn_cells', 'class_net', 'box_net', 'Reshape', 'concat')
|
||||
|
||||
addConstNode('scale_w', [2.0], graph_def)
|
||||
addConstNode('scale_h', [2.0], graph_def)
|
||||
nodesToKeep += ['scale_w', 'scale_h']
|
||||
|
||||
for node in graph_def.node:
|
||||
if re.match('efficientnet-(.*)/blocks_\d+/se/mul_1', node.name):
|
||||
node.input[0], node.input[1] = node.input[1], node.input[0]
|
||||
|
||||
if re.match('fpn_cells/cell_\d+/fnode\d+/resample(.*)/nearest_upsampling/Reshape_1$', node.name):
|
||||
node.op = 'ResizeNearestNeighbor'
|
||||
node.input[1] = 'scale_w'
|
||||
node.input.append('scale_h')
|
||||
|
||||
for inpNode in graph_def.node:
|
||||
if inpNode.name == node.name[:node.name.rfind('_')]:
|
||||
node.input[0] = inpNode.input[0]
|
||||
|
||||
if re.match('box_net/box-predict(_\d)*/separable_conv2d$', node.name):
|
||||
node.addAttr('loc_pred_transposed', True)
|
||||
|
||||
# Replace RealDiv to Mul with inversed scale for compatibility
|
||||
if node.op == 'RealDiv':
|
||||
for inpNode in graph_def.node:
|
||||
if inpNode.name != node.input[1] or not 'value' in inpNode.attr:
|
||||
continue
|
||||
|
||||
tensor = inpNode.attr['value']['tensor'][0]
|
||||
if not 'float_val' in tensor:
|
||||
continue
|
||||
scale = float(inpNode.attr['value']['tensor'][0]['float_val'][0])
|
||||
|
||||
addConstNode(inpNode.name + '/inv', [1.0 / scale], graph_def)
|
||||
nodesToKeep.append(inpNode.name + '/inv')
|
||||
node.input[1] = inpNode.name + '/inv'
|
||||
node.op = 'Mul'
|
||||
break
|
||||
|
||||
|
||||
def to_remove(name, op):
|
||||
if name in nodesToKeep:
|
||||
return False
|
||||
return op == 'Const' or not name.startswith(scopesToKeep)
|
||||
|
||||
removeUnusedNodesAndAttrs(to_remove, graph_def)
|
||||
|
||||
# Attach unconnected preprocessing
|
||||
assert(graph_def.node[1].name == 'truediv' and graph_def.node[1].op == 'RealDiv')
|
||||
graph_def.node[1].input.insert(0, 'image_arrays')
|
||||
graph_def.node[2].input.insert(0, 'truediv')
|
||||
|
||||
priors_generator = AnchorGenerator(min_level, aspect_ratios, num_scales, anchor_scale)
|
||||
priorBoxes = []
|
||||
for i in range(5):
|
||||
inpName = ''
|
||||
for node in graph_def.node:
|
||||
if node.name == 'Reshape_%d' % (i * 2 + 1):
|
||||
inpName = node.input[0]
|
||||
break
|
||||
|
||||
priorBox = NodeDef()
|
||||
priorBox.name = 'PriorBox_%d' % i
|
||||
priorBox.op = 'PriorBox'
|
||||
priorBox.input.append(inpName)
|
||||
priorBox.input.append(graph_def.node[0].name) # image_tensor
|
||||
|
||||
priorBox.addAttr('flip', False)
|
||||
priorBox.addAttr('clip', False)
|
||||
|
||||
widths, heights = priors_generator.get(i)
|
||||
|
||||
priorBox.addAttr('width', widths)
|
||||
priorBox.addAttr('height', heights)
|
||||
priorBox.addAttr('variance', [1.0, 1.0, 1.0, 1.0])
|
||||
|
||||
graph_def.node.extend([priorBox])
|
||||
priorBoxes.append(priorBox.name)
|
||||
|
||||
addConstNode('concat/axis_flatten', [-1], graph_def)
|
||||
|
||||
def addConcatNode(name, inputs, axisNodeName):
|
||||
concat = NodeDef()
|
||||
concat.name = name
|
||||
concat.op = 'ConcatV2'
|
||||
for inp in inputs:
|
||||
concat.input.append(inp)
|
||||
concat.input.append(axisNodeName)
|
||||
graph_def.node.extend([concat])
|
||||
|
||||
addConcatNode('PriorBox/concat', priorBoxes, 'concat/axis_flatten')
|
||||
|
||||
sigmoid = NodeDef()
|
||||
sigmoid.name = 'concat/sigmoid'
|
||||
sigmoid.op = 'Sigmoid'
|
||||
sigmoid.input.append('concat')
|
||||
graph_def.node.extend([sigmoid])
|
||||
|
||||
addFlatten(sigmoid.name, sigmoid.name + '/Flatten', graph_def)
|
||||
addFlatten('concat_1', 'concat_1/Flatten', graph_def)
|
||||
|
||||
detectionOut = NodeDef()
|
||||
detectionOut.name = 'detection_out'
|
||||
detectionOut.op = 'DetectionOutput'
|
||||
|
||||
detectionOut.input.append('concat_1/Flatten')
|
||||
detectionOut.input.append(sigmoid.name + '/Flatten')
|
||||
detectionOut.input.append('PriorBox/concat')
|
||||
|
||||
detectionOut.addAttr('num_classes', num_classes)
|
||||
detectionOut.addAttr('share_location', True)
|
||||
detectionOut.addAttr('background_label_id', num_classes + 1)
|
||||
detectionOut.addAttr('nms_threshold', 0.6)
|
||||
detectionOut.addAttr('confidence_threshold', 0.2)
|
||||
detectionOut.addAttr('top_k', 100)
|
||||
detectionOut.addAttr('keep_top_k', 100)
|
||||
detectionOut.addAttr('code_type', "CENTER_SIZE")
|
||||
graph_def.node.extend([detectionOut])
|
||||
|
||||
graph_def.node[0].attr['shape'] = {
|
||||
'shape': {
|
||||
'dim': [
|
||||
{'size': -1},
|
||||
{'size': image_height},
|
||||
{'size': image_width},
|
||||
{'size': 3}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
while True:
|
||||
unconnectedNodes = getUnconnectedNodes()
|
||||
unconnectedNodes.remove(detectionOut.name)
|
||||
if not unconnectedNodes:
|
||||
break
|
||||
|
||||
for name in unconnectedNodes:
|
||||
for i in range(len(graph_def.node)):
|
||||
if graph_def.node[i].name == name:
|
||||
del graph_def.node[i]
|
||||
break
|
||||
|
||||
# Save as text
|
||||
graph_def.save(outputPath)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Run this script to get a text graph of '
|
||||
'SSD model from TensorFlow Object Detection API. '
|
||||
'Then pass it with .pb file to cv::dnn::readNetFromTensorflow function.')
|
||||
parser.add_argument('--input', required=True, help='Path to frozen TensorFlow graph.')
|
||||
parser.add_argument('--output', required=True, help='Path to output text graph.')
|
||||
parser.add_argument('--min_level', default=3, type=int, help='Parameter from training config')
|
||||
parser.add_argument('--num_scales', default=3, type=int, help='Parameter from training config')
|
||||
parser.add_argument('--anchor_scale', default=4.0, type=float, help='Parameter from training config')
|
||||
parser.add_argument('--aspect_ratios', default=[1.0, 1.0, 1.4, 0.7, 0.7, 1.4],
|
||||
nargs='+', type=float, help='Parameter from training config')
|
||||
parser.add_argument('--num_classes', default=90, type=int, help='Number of classes to detect')
|
||||
parser.add_argument('--width', default=512, type=int, help='Network input width')
|
||||
parser.add_argument('--height', default=512, type=int, help='Network input height')
|
||||
args = parser.parse_args()
|
||||
|
||||
ar = args.aspect_ratios
|
||||
assert(len(ar) % 2 == 0)
|
||||
ar = list(zip(ar[::2], ar[1::2]))
|
||||
|
||||
createGraph(args.input, args.output, args.min_level, ar, args.num_scales,
|
||||
args.anchor_scale, args.num_classes, args.width, args.height)
|
||||
@@ -0,0 +1,299 @@
|
||||
import argparse
|
||||
import numpy as np
|
||||
from tf_text_graph_common import *
|
||||
|
||||
|
||||
def createFasterRCNNGraph(modelPath, configPath, outputPath):
|
||||
scopesToKeep = ('FirstStageFeatureExtractor', 'Conv',
|
||||
'FirstStageBoxPredictor/BoxEncodingPredictor',
|
||||
'FirstStageBoxPredictor/ClassPredictor',
|
||||
'CropAndResize',
|
||||
'MaxPool2D',
|
||||
'SecondStageFeatureExtractor',
|
||||
'SecondStageBoxPredictor',
|
||||
'Preprocessor/sub',
|
||||
'Preprocessor/mul',
|
||||
'image_tensor')
|
||||
|
||||
scopesToIgnore = ('FirstStageFeatureExtractor/Assert',
|
||||
'FirstStageFeatureExtractor/Shape',
|
||||
'FirstStageFeatureExtractor/strided_slice',
|
||||
'FirstStageFeatureExtractor/GreaterEqual',
|
||||
'FirstStageFeatureExtractor/LogicalAnd')
|
||||
|
||||
# Load a config file.
|
||||
config = readTextMessage(configPath)
|
||||
config = config['model'][0]['faster_rcnn'][0]
|
||||
num_classes = int(config['num_classes'][0])
|
||||
|
||||
grid_anchor_generator = config['first_stage_anchor_generator'][0]['grid_anchor_generator'][0]
|
||||
scales = [float(s) for s in grid_anchor_generator['scales']]
|
||||
aspect_ratios = [float(ar) for ar in grid_anchor_generator['aspect_ratios']]
|
||||
width_stride = float(grid_anchor_generator['width_stride'][0])
|
||||
height_stride = float(grid_anchor_generator['height_stride'][0])
|
||||
|
||||
feature_extractor = config['feature_extractor'][0]
|
||||
if 'type' in feature_extractor and feature_extractor['type'][0] == 'faster_rcnn_nas':
|
||||
features_stride = 16.0
|
||||
else:
|
||||
features_stride = float(feature_extractor['first_stage_features_stride'][0])
|
||||
|
||||
first_stage_nms_iou_threshold = float(config['first_stage_nms_iou_threshold'][0])
|
||||
first_stage_max_proposals = int(config['first_stage_max_proposals'][0])
|
||||
|
||||
print('Number of classes: %d' % num_classes)
|
||||
print('Scales: %s' % str(scales))
|
||||
print('Aspect ratios: %s' % str(aspect_ratios))
|
||||
print('Width stride: %f' % width_stride)
|
||||
print('Height stride: %f' % height_stride)
|
||||
print('Features stride: %f' % features_stride)
|
||||
|
||||
# Read the graph.
|
||||
writeTextGraph(modelPath, outputPath, ['num_detections', 'detection_scores', 'detection_boxes', 'detection_classes'])
|
||||
graph_def = parseTextGraph(outputPath)
|
||||
|
||||
removeIdentity(graph_def)
|
||||
|
||||
nodesToKeep = []
|
||||
def to_remove(name, op):
|
||||
if name in nodesToKeep:
|
||||
return False
|
||||
return op == 'Const' or name.startswith(scopesToIgnore) or not name.startswith(scopesToKeep) or \
|
||||
(name.startswith('CropAndResize') and op != 'CropAndResize')
|
||||
|
||||
# Fuse atrous convolutions (with dilations).
|
||||
nodesMap = {node.name: node for node in graph_def.node}
|
||||
for node in reversed(graph_def.node):
|
||||
if node.op == 'BatchToSpaceND':
|
||||
del node.input[2]
|
||||
conv = nodesMap[node.input[0]]
|
||||
spaceToBatchND = nodesMap[conv.input[0]]
|
||||
|
||||
# Extract paddings
|
||||
stridedSlice = nodesMap[spaceToBatchND.input[2]]
|
||||
assert(stridedSlice.op == 'StridedSlice')
|
||||
pack = nodesMap[stridedSlice.input[0]]
|
||||
assert(pack.op == 'Pack')
|
||||
|
||||
padNodeH = nodesMap[nodesMap[pack.input[0]].input[0]]
|
||||
padNodeW = nodesMap[nodesMap[pack.input[1]].input[0]]
|
||||
padH = int(padNodeH.attr['value']['tensor'][0]['int_val'][0])
|
||||
padW = int(padNodeW.attr['value']['tensor'][0]['int_val'][0])
|
||||
|
||||
paddingsNode = NodeDef()
|
||||
paddingsNode.name = conv.name + '/paddings'
|
||||
paddingsNode.op = 'Const'
|
||||
paddingsNode.addAttr('value', [padH, padH, padW, padW])
|
||||
graph_def.node.insert(graph_def.node.index(spaceToBatchND), paddingsNode)
|
||||
nodesToKeep.append(paddingsNode.name)
|
||||
|
||||
spaceToBatchND.input[2] = paddingsNode.name
|
||||
|
||||
|
||||
removeUnusedNodesAndAttrs(to_remove, graph_def)
|
||||
|
||||
|
||||
# Connect input node to the first layer
|
||||
assert(graph_def.node[0].op == 'Placeholder')
|
||||
graph_def.node[1].input.insert(0, graph_def.node[0].name)
|
||||
|
||||
# Temporarily remove top nodes.
|
||||
topNodes = []
|
||||
while True:
|
||||
node = graph_def.node.pop()
|
||||
topNodes.append(node)
|
||||
if node.op == 'CropAndResize':
|
||||
break
|
||||
|
||||
addReshape('FirstStageBoxPredictor/ClassPredictor/BiasAdd',
|
||||
'FirstStageBoxPredictor/ClassPredictor/reshape_1', [0, -1, 2], graph_def)
|
||||
|
||||
addSoftMax('FirstStageBoxPredictor/ClassPredictor/reshape_1',
|
||||
'FirstStageBoxPredictor/ClassPredictor/softmax', graph_def) # Compare with Reshape_4
|
||||
|
||||
addFlatten('FirstStageBoxPredictor/ClassPredictor/softmax',
|
||||
'FirstStageBoxPredictor/ClassPredictor/softmax/flatten', graph_def)
|
||||
|
||||
# Compare with FirstStageBoxPredictor/BoxEncodingPredictor/BiasAdd
|
||||
addFlatten('FirstStageBoxPredictor/BoxEncodingPredictor/BiasAdd',
|
||||
'FirstStageBoxPredictor/BoxEncodingPredictor/flatten', graph_def)
|
||||
|
||||
proposals = NodeDef()
|
||||
proposals.name = 'proposals' # Compare with ClipToWindow/Gather/Gather (NOTE: normalized)
|
||||
proposals.op = 'PriorBox'
|
||||
proposals.input.append('FirstStageBoxPredictor/BoxEncodingPredictor/BiasAdd')
|
||||
proposals.input.append(graph_def.node[0].name) # image_tensor
|
||||
|
||||
proposals.addAttr('flip', False)
|
||||
proposals.addAttr('clip', True)
|
||||
proposals.addAttr('step', features_stride)
|
||||
proposals.addAttr('offset', 0.0)
|
||||
proposals.addAttr('variance', [0.1, 0.1, 0.2, 0.2])
|
||||
|
||||
widths = []
|
||||
heights = []
|
||||
for a in aspect_ratios:
|
||||
for s in scales:
|
||||
ar = np.sqrt(a)
|
||||
heights.append((height_stride**2) * s / ar)
|
||||
widths.append((width_stride**2) * s * ar)
|
||||
|
||||
proposals.addAttr('width', widths)
|
||||
proposals.addAttr('height', heights)
|
||||
|
||||
graph_def.node.extend([proposals])
|
||||
|
||||
# Compare with Reshape_5
|
||||
detectionOut = NodeDef()
|
||||
detectionOut.name = 'detection_out'
|
||||
detectionOut.op = 'DetectionOutput'
|
||||
|
||||
detectionOut.input.append('FirstStageBoxPredictor/BoxEncodingPredictor/flatten')
|
||||
detectionOut.input.append('FirstStageBoxPredictor/ClassPredictor/softmax/flatten')
|
||||
detectionOut.input.append('proposals')
|
||||
|
||||
detectionOut.addAttr('num_classes', 2)
|
||||
detectionOut.addAttr('share_location', True)
|
||||
detectionOut.addAttr('background_label_id', 0)
|
||||
detectionOut.addAttr('nms_threshold', first_stage_nms_iou_threshold)
|
||||
detectionOut.addAttr('top_k', 6000)
|
||||
detectionOut.addAttr('code_type', "CENTER_SIZE")
|
||||
detectionOut.addAttr('keep_top_k', first_stage_max_proposals)
|
||||
detectionOut.addAttr('clip', False)
|
||||
|
||||
graph_def.node.extend([detectionOut])
|
||||
|
||||
addConstNode('clip_by_value/lower', [0.0], graph_def)
|
||||
addConstNode('clip_by_value/upper', [1.0], graph_def)
|
||||
|
||||
clipByValueNode = NodeDef()
|
||||
clipByValueNode.name = 'detection_out/clip_by_value'
|
||||
clipByValueNode.op = 'ClipByValue'
|
||||
clipByValueNode.input.append('detection_out')
|
||||
clipByValueNode.input.append('clip_by_value/lower')
|
||||
clipByValueNode.input.append('clip_by_value/upper')
|
||||
graph_def.node.extend([clipByValueNode])
|
||||
|
||||
# Save as text.
|
||||
for node in reversed(topNodes):
|
||||
graph_def.node.extend([node])
|
||||
|
||||
addSoftMax('SecondStageBoxPredictor/Reshape_1', 'SecondStageBoxPredictor/Reshape_1/softmax', graph_def)
|
||||
|
||||
addSlice('SecondStageBoxPredictor/Reshape_1/softmax',
|
||||
'SecondStageBoxPredictor/Reshape_1/slice',
|
||||
[0, 0, 1], [-1, -1, -1], graph_def)
|
||||
|
||||
addReshape('SecondStageBoxPredictor/Reshape_1/slice',
|
||||
'SecondStageBoxPredictor/Reshape_1/Reshape', [1, -1], graph_def)
|
||||
|
||||
# Replace Flatten subgraph onto a single node.
|
||||
cropAndResizeNodeName = ''
|
||||
for i in reversed(range(len(graph_def.node))):
|
||||
if graph_def.node[i].op == 'CropAndResize':
|
||||
graph_def.node[i].input.insert(1, 'detection_out/clip_by_value')
|
||||
cropAndResizeNodeName = graph_def.node[i].name
|
||||
|
||||
if graph_def.node[i].name == 'SecondStageBoxPredictor/Reshape':
|
||||
addConstNode('SecondStageBoxPredictor/Reshape/shape2', [1, -1, 4], graph_def)
|
||||
|
||||
graph_def.node[i].input.pop()
|
||||
graph_def.node[i].input.append('SecondStageBoxPredictor/Reshape/shape2')
|
||||
|
||||
if graph_def.node[i].name in ['SecondStageBoxPredictor/Flatten/flatten/Shape',
|
||||
'SecondStageBoxPredictor/Flatten/flatten/strided_slice',
|
||||
'SecondStageBoxPredictor/Flatten/flatten/Reshape/shape',
|
||||
'SecondStageBoxPredictor/Flatten_1/flatten/Shape',
|
||||
'SecondStageBoxPredictor/Flatten_1/flatten/strided_slice',
|
||||
'SecondStageBoxPredictor/Flatten_1/flatten/Reshape/shape']:
|
||||
del graph_def.node[i]
|
||||
|
||||
for node in graph_def.node:
|
||||
if node.name == 'SecondStageBoxPredictor/Flatten/flatten/Reshape' or \
|
||||
node.name == 'SecondStageBoxPredictor/Flatten_1/flatten/Reshape':
|
||||
node.op = 'Flatten'
|
||||
node.input.pop()
|
||||
|
||||
if node.name in ['FirstStageBoxPredictor/BoxEncodingPredictor/Conv2D',
|
||||
'SecondStageBoxPredictor/BoxEncodingPredictor/MatMul']:
|
||||
node.addAttr('loc_pred_transposed', True)
|
||||
|
||||
if node.name.startswith('MaxPool2D'):
|
||||
assert(node.op == 'MaxPool')
|
||||
assert(cropAndResizeNodeName)
|
||||
node.input = [cropAndResizeNodeName]
|
||||
|
||||
################################################################################
|
||||
### Postprocessing
|
||||
################################################################################
|
||||
addSlice('detection_out/clip_by_value', 'detection_out/slice', [0, 0, 0, 3], [-1, -1, -1, 4], graph_def)
|
||||
|
||||
variance = NodeDef()
|
||||
variance.name = 'proposals/variance'
|
||||
variance.op = 'Const'
|
||||
variance.addAttr('value', [0.1, 0.1, 0.2, 0.2])
|
||||
graph_def.node.extend([variance])
|
||||
|
||||
varianceEncoder = NodeDef()
|
||||
varianceEncoder.name = 'variance_encoded'
|
||||
varianceEncoder.op = 'Mul'
|
||||
varianceEncoder.input.append('SecondStageBoxPredictor/Reshape')
|
||||
varianceEncoder.input.append(variance.name)
|
||||
varianceEncoder.addAttr('axis', 2)
|
||||
graph_def.node.extend([varianceEncoder])
|
||||
|
||||
addReshape('detection_out/slice', 'detection_out/slice/reshape', [1, 1, -1], graph_def)
|
||||
addFlatten('variance_encoded', 'variance_encoded/flatten', graph_def)
|
||||
|
||||
detectionOut = NodeDef()
|
||||
detectionOut.name = 'detection_out_final'
|
||||
detectionOut.op = 'DetectionOutput'
|
||||
|
||||
detectionOut.input.append('variance_encoded/flatten')
|
||||
detectionOut.input.append('SecondStageBoxPredictor/Reshape_1/Reshape')
|
||||
detectionOut.input.append('detection_out/slice/reshape')
|
||||
|
||||
detectionOut.addAttr('num_classes', num_classes)
|
||||
detectionOut.addAttr('share_location', False)
|
||||
detectionOut.addAttr('background_label_id', num_classes + 1)
|
||||
detectionOut.addAttr('nms_threshold', 0.6)
|
||||
detectionOut.addAttr('code_type', "CENTER_SIZE")
|
||||
detectionOut.addAttr('keep_top_k', 100)
|
||||
detectionOut.addAttr('clip', True)
|
||||
detectionOut.addAttr('variance_encoded_in_target', True)
|
||||
graph_def.node.extend([detectionOut])
|
||||
|
||||
def getUnconnectedNodes():
|
||||
unconnected = [node.name for node in graph_def.node]
|
||||
for node in graph_def.node:
|
||||
for inp in node.input:
|
||||
if inp in unconnected:
|
||||
unconnected.remove(inp)
|
||||
return unconnected
|
||||
|
||||
while True:
|
||||
unconnectedNodes = getUnconnectedNodes()
|
||||
unconnectedNodes.remove(detectionOut.name)
|
||||
if not unconnectedNodes:
|
||||
break
|
||||
|
||||
for name in unconnectedNodes:
|
||||
for i in range(len(graph_def.node)):
|
||||
if graph_def.node[i].name == name:
|
||||
del graph_def.node[i]
|
||||
break
|
||||
|
||||
# Save as text.
|
||||
graph_def.save(outputPath)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Run this script to get a text graph of '
|
||||
'Faster-RCNN model from TensorFlow Object Detection API. '
|
||||
'Then pass it with .pb file to cv::dnn::readNetFromTensorflow function.')
|
||||
parser.add_argument('--input', required=True, help='Path to frozen TensorFlow graph.')
|
||||
parser.add_argument('--output', required=True, help='Path to output text graph.')
|
||||
parser.add_argument('--config', required=True, help='Path to a *.config file is used for training.')
|
||||
args = parser.parse_args()
|
||||
|
||||
createFasterRCNNGraph(args.input, args.config, args.output)
|
||||
@@ -0,0 +1,297 @@
|
||||
import argparse
|
||||
import numpy as np
|
||||
from tf_text_graph_common import *
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run this script to get a text graph of '
|
||||
'Mask-RCNN model from TensorFlow Object Detection API. '
|
||||
'Then pass it with .pb file to cv::dnn::readNetFromTensorflow function.')
|
||||
parser.add_argument('--input', required=True, help='Path to frozen TensorFlow graph.')
|
||||
parser.add_argument('--output', required=True, help='Path to output text graph.')
|
||||
parser.add_argument('--config', required=True, help='Path to a *.config file is used for training.')
|
||||
args = parser.parse_args()
|
||||
|
||||
scopesToKeep = ('FirstStageFeatureExtractor', 'Conv',
|
||||
'FirstStageBoxPredictor/BoxEncodingPredictor',
|
||||
'FirstStageBoxPredictor/ClassPredictor',
|
||||
'CropAndResize',
|
||||
'MaxPool2D',
|
||||
'SecondStageFeatureExtractor',
|
||||
'SecondStageBoxPredictor',
|
||||
'Preprocessor/sub',
|
||||
'Preprocessor/mul',
|
||||
'image_tensor')
|
||||
|
||||
scopesToIgnore = ('FirstStageFeatureExtractor/Assert',
|
||||
'FirstStageFeatureExtractor/Shape',
|
||||
'FirstStageFeatureExtractor/strided_slice',
|
||||
'FirstStageFeatureExtractor/GreaterEqual',
|
||||
'FirstStageFeatureExtractor/LogicalAnd',
|
||||
'Conv/required_space_to_batch_paddings')
|
||||
|
||||
# Load a config file.
|
||||
config = readTextMessage(args.config)
|
||||
config = config['model'][0]['faster_rcnn'][0]
|
||||
num_classes = int(config['num_classes'][0])
|
||||
|
||||
grid_anchor_generator = config['first_stage_anchor_generator'][0]['grid_anchor_generator'][0]
|
||||
scales = [float(s) for s in grid_anchor_generator['scales']]
|
||||
aspect_ratios = [float(ar) for ar in grid_anchor_generator['aspect_ratios']]
|
||||
width_stride = float(grid_anchor_generator['width_stride'][0])
|
||||
height_stride = float(grid_anchor_generator['height_stride'][0])
|
||||
features_stride = float(config['feature_extractor'][0]['first_stage_features_stride'][0])
|
||||
first_stage_nms_iou_threshold = float(config['first_stage_nms_iou_threshold'][0])
|
||||
first_stage_max_proposals = int(config['first_stage_max_proposals'][0])
|
||||
|
||||
print('Number of classes: %d' % num_classes)
|
||||
print('Scales: %s' % str(scales))
|
||||
print('Aspect ratios: %s' % str(aspect_ratios))
|
||||
print('Width stride: %f' % width_stride)
|
||||
print('Height stride: %f' % height_stride)
|
||||
print('Features stride: %f' % features_stride)
|
||||
|
||||
# Read the graph.
|
||||
writeTextGraph(args.input, args.output, ['num_detections', 'detection_scores', 'detection_boxes', 'detection_classes', 'detection_masks'])
|
||||
graph_def = parseTextGraph(args.output)
|
||||
|
||||
removeIdentity(graph_def)
|
||||
|
||||
nodesToKeep = []
|
||||
def to_remove(name, op):
|
||||
if name in nodesToKeep:
|
||||
return False
|
||||
return op == 'Const' or name.startswith(scopesToIgnore) or not name.startswith(scopesToKeep) or \
|
||||
(name.startswith('CropAndResize') and op != 'CropAndResize')
|
||||
|
||||
# Fuse atrous convolutions (with dilations).
|
||||
nodesMap = {node.name: node for node in graph_def.node}
|
||||
for node in reversed(graph_def.node):
|
||||
if node.op == 'BatchToSpaceND':
|
||||
del node.input[2]
|
||||
conv = nodesMap[node.input[0]]
|
||||
spaceToBatchND = nodesMap[conv.input[0]]
|
||||
|
||||
paddingsNode = NodeDef()
|
||||
paddingsNode.name = conv.name + '/paddings'
|
||||
paddingsNode.op = 'Const'
|
||||
paddingsNode.addAttr('value', [2, 2, 2, 2])
|
||||
graph_def.node.insert(graph_def.node.index(spaceToBatchND), paddingsNode)
|
||||
nodesToKeep.append(paddingsNode.name)
|
||||
|
||||
spaceToBatchND.input[2] = paddingsNode.name
|
||||
|
||||
removeUnusedNodesAndAttrs(to_remove, graph_def)
|
||||
|
||||
|
||||
# Connect input node to the first layer
|
||||
assert(graph_def.node[0].op == 'Placeholder')
|
||||
graph_def.node[1].input.insert(0, graph_def.node[0].name)
|
||||
|
||||
# Temporarily remove top nodes.
|
||||
topNodes = []
|
||||
numCropAndResize = 0
|
||||
while True:
|
||||
node = graph_def.node.pop()
|
||||
topNodes.append(node)
|
||||
if node.op == 'CropAndResize':
|
||||
numCropAndResize += 1
|
||||
if numCropAndResize == 2:
|
||||
break
|
||||
|
||||
addReshape('FirstStageBoxPredictor/ClassPredictor/BiasAdd',
|
||||
'FirstStageBoxPredictor/ClassPredictor/reshape_1', [0, -1, 2], graph_def)
|
||||
|
||||
addSoftMax('FirstStageBoxPredictor/ClassPredictor/reshape_1',
|
||||
'FirstStageBoxPredictor/ClassPredictor/softmax', graph_def) # Compare with Reshape_4
|
||||
|
||||
addFlatten('FirstStageBoxPredictor/ClassPredictor/softmax',
|
||||
'FirstStageBoxPredictor/ClassPredictor/softmax/flatten', graph_def)
|
||||
|
||||
# Compare with FirstStageBoxPredictor/BoxEncodingPredictor/BiasAdd
|
||||
addFlatten('FirstStageBoxPredictor/BoxEncodingPredictor/BiasAdd',
|
||||
'FirstStageBoxPredictor/BoxEncodingPredictor/flatten', graph_def)
|
||||
|
||||
proposals = NodeDef()
|
||||
proposals.name = 'proposals' # Compare with ClipToWindow/Gather/Gather (NOTE: normalized)
|
||||
proposals.op = 'PriorBox'
|
||||
proposals.input.append('FirstStageBoxPredictor/BoxEncodingPredictor/BiasAdd')
|
||||
proposals.input.append(graph_def.node[0].name) # image_tensor
|
||||
|
||||
proposals.addAttr('flip', False)
|
||||
proposals.addAttr('clip', True)
|
||||
proposals.addAttr('step', features_stride)
|
||||
proposals.addAttr('offset', 0.0)
|
||||
proposals.addAttr('variance', [0.1, 0.1, 0.2, 0.2])
|
||||
|
||||
widths = []
|
||||
heights = []
|
||||
for a in aspect_ratios:
|
||||
for s in scales:
|
||||
ar = np.sqrt(a)
|
||||
heights.append((height_stride**2) * s / ar)
|
||||
widths.append((width_stride**2) * s * ar)
|
||||
|
||||
proposals.addAttr('width', widths)
|
||||
proposals.addAttr('height', heights)
|
||||
|
||||
graph_def.node.extend([proposals])
|
||||
|
||||
# Compare with Reshape_5
|
||||
detectionOut = NodeDef()
|
||||
detectionOut.name = 'detection_out'
|
||||
detectionOut.op = 'DetectionOutput'
|
||||
|
||||
detectionOut.input.append('FirstStageBoxPredictor/BoxEncodingPredictor/flatten')
|
||||
detectionOut.input.append('FirstStageBoxPredictor/ClassPredictor/softmax/flatten')
|
||||
detectionOut.input.append('proposals')
|
||||
|
||||
detectionOut.addAttr('num_classes', 2)
|
||||
detectionOut.addAttr('share_location', True)
|
||||
detectionOut.addAttr('background_label_id', 0)
|
||||
detectionOut.addAttr('nms_threshold', first_stage_nms_iou_threshold)
|
||||
detectionOut.addAttr('top_k', 6000)
|
||||
detectionOut.addAttr('code_type', "CENTER_SIZE")
|
||||
detectionOut.addAttr('keep_top_k', first_stage_max_proposals)
|
||||
detectionOut.addAttr('clip', True)
|
||||
|
||||
graph_def.node.extend([detectionOut])
|
||||
|
||||
# Save as text.
|
||||
cropAndResizeNodesNames = []
|
||||
for node in reversed(topNodes):
|
||||
if node.op != 'CropAndResize':
|
||||
graph_def.node.extend([node])
|
||||
topNodes.pop()
|
||||
else:
|
||||
cropAndResizeNodesNames.append(node.name)
|
||||
if numCropAndResize == 1:
|
||||
break
|
||||
else:
|
||||
graph_def.node.extend([node])
|
||||
topNodes.pop()
|
||||
numCropAndResize -= 1
|
||||
|
||||
addSoftMax('SecondStageBoxPredictor/Reshape_1', 'SecondStageBoxPredictor/Reshape_1/softmax', graph_def)
|
||||
|
||||
addSlice('SecondStageBoxPredictor/Reshape_1/softmax',
|
||||
'SecondStageBoxPredictor/Reshape_1/slice',
|
||||
[0, 0, 1], [-1, -1, -1], graph_def)
|
||||
|
||||
addReshape('SecondStageBoxPredictor/Reshape_1/slice',
|
||||
'SecondStageBoxPredictor/Reshape_1/Reshape', [1, -1], graph_def)
|
||||
|
||||
# Replace Flatten subgraph onto a single node.
|
||||
for i in reversed(range(len(graph_def.node))):
|
||||
if graph_def.node[i].op == 'CropAndResize':
|
||||
graph_def.node[i].input.insert(1, 'detection_out')
|
||||
|
||||
if graph_def.node[i].name == 'SecondStageBoxPredictor/Reshape':
|
||||
addConstNode('SecondStageBoxPredictor/Reshape/shape2', [1, -1, 4], graph_def)
|
||||
|
||||
graph_def.node[i].input.pop()
|
||||
graph_def.node[i].input.append('SecondStageBoxPredictor/Reshape/shape2')
|
||||
|
||||
if graph_def.node[i].name in ['SecondStageBoxPredictor/Flatten/flatten/Shape',
|
||||
'SecondStageBoxPredictor/Flatten/flatten/strided_slice',
|
||||
'SecondStageBoxPredictor/Flatten/flatten/Reshape/shape',
|
||||
'SecondStageBoxPredictor/Flatten_1/flatten/Shape',
|
||||
'SecondStageBoxPredictor/Flatten_1/flatten/strided_slice',
|
||||
'SecondStageBoxPredictor/Flatten_1/flatten/Reshape/shape']:
|
||||
del graph_def.node[i]
|
||||
|
||||
for node in graph_def.node:
|
||||
if node.name == 'SecondStageBoxPredictor/Flatten/flatten/Reshape' or \
|
||||
node.name == 'SecondStageBoxPredictor/Flatten_1/flatten/Reshape':
|
||||
node.op = 'Flatten'
|
||||
node.input.pop()
|
||||
|
||||
if node.name in ['FirstStageBoxPredictor/BoxEncodingPredictor/Conv2D',
|
||||
'SecondStageBoxPredictor/BoxEncodingPredictor/MatMul']:
|
||||
node.addAttr('loc_pred_transposed', True)
|
||||
|
||||
if node.name.startswith('MaxPool2D'):
|
||||
assert(node.op == 'MaxPool')
|
||||
assert(len(cropAndResizeNodesNames) == 2)
|
||||
node.input = [cropAndResizeNodesNames[0]]
|
||||
del cropAndResizeNodesNames[0]
|
||||
|
||||
################################################################################
|
||||
### Postprocessing
|
||||
################################################################################
|
||||
addSlice('detection_out', 'detection_out/slice', [0, 0, 0, 3], [-1, -1, -1, 4], graph_def)
|
||||
|
||||
variance = NodeDef()
|
||||
variance.name = 'proposals/variance'
|
||||
variance.op = 'Const'
|
||||
variance.addAttr('value', [0.1, 0.1, 0.2, 0.2])
|
||||
graph_def.node.extend([variance])
|
||||
|
||||
varianceEncoder = NodeDef()
|
||||
varianceEncoder.name = 'variance_encoded'
|
||||
varianceEncoder.op = 'Mul'
|
||||
varianceEncoder.input.append('SecondStageBoxPredictor/Reshape')
|
||||
varianceEncoder.input.append(variance.name)
|
||||
varianceEncoder.addAttr('axis', 2)
|
||||
graph_def.node.extend([varianceEncoder])
|
||||
|
||||
addReshape('detection_out/slice', 'detection_out/slice/reshape', [1, 1, -1], graph_def)
|
||||
addFlatten('variance_encoded', 'variance_encoded/flatten', graph_def)
|
||||
|
||||
detectionOut = NodeDef()
|
||||
detectionOut.name = 'detection_out_final'
|
||||
detectionOut.op = 'DetectionOutput'
|
||||
|
||||
detectionOut.input.append('variance_encoded/flatten')
|
||||
detectionOut.input.append('SecondStageBoxPredictor/Reshape_1/Reshape')
|
||||
detectionOut.input.append('detection_out/slice/reshape')
|
||||
|
||||
detectionOut.addAttr('num_classes', num_classes)
|
||||
detectionOut.addAttr('share_location', False)
|
||||
detectionOut.addAttr('background_label_id', num_classes + 1)
|
||||
detectionOut.addAttr('nms_threshold', 0.6)
|
||||
detectionOut.addAttr('code_type', "CENTER_SIZE")
|
||||
detectionOut.addAttr('keep_top_k',100)
|
||||
detectionOut.addAttr('clip', True)
|
||||
detectionOut.addAttr('variance_encoded_in_target', True)
|
||||
detectionOut.addAttr('confidence_threshold', 0.3)
|
||||
detectionOut.addAttr('group_by_classes', False)
|
||||
graph_def.node.extend([detectionOut])
|
||||
|
||||
for node in reversed(topNodes):
|
||||
graph_def.node.extend([node])
|
||||
|
||||
if node.name.startswith('MaxPool2D'):
|
||||
assert(node.op == 'MaxPool')
|
||||
assert(len(cropAndResizeNodesNames) == 1)
|
||||
node.input = [cropAndResizeNodesNames[0]]
|
||||
|
||||
for i in reversed(range(len(graph_def.node))):
|
||||
if graph_def.node[i].op == 'CropAndResize':
|
||||
graph_def.node[i].input.insert(1, 'detection_out_final')
|
||||
break
|
||||
|
||||
graph_def.node[-1].name = 'detection_masks'
|
||||
graph_def.node[-1].op = 'Sigmoid'
|
||||
graph_def.node[-1].input.pop()
|
||||
|
||||
def getUnconnectedNodes():
|
||||
unconnected = [node.name for node in graph_def.node]
|
||||
for node in graph_def.node:
|
||||
for inp in node.input:
|
||||
if inp in unconnected:
|
||||
unconnected.remove(inp)
|
||||
return unconnected
|
||||
|
||||
while True:
|
||||
unconnectedNodes = getUnconnectedNodes()
|
||||
unconnectedNodes.remove(graph_def.node[-1].name)
|
||||
if not unconnectedNodes:
|
||||
break
|
||||
|
||||
for name in unconnectedNodes:
|
||||
for i in range(len(graph_def.node)):
|
||||
if graph_def.node[i].name == name:
|
||||
del graph_def.node[i]
|
||||
break
|
||||
|
||||
# Save as text.
|
||||
graph_def.save(args.output)
|
||||
@@ -0,0 +1,413 @@
|
||||
# This file is a part of OpenCV project.
|
||||
# It is a 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) 2018, Intel Corporation, all rights reserved.
|
||||
# Third party copyrights are property of their respective owners.
|
||||
#
|
||||
# Use this script to get the text graph representation (.pbtxt) of SSD-based
|
||||
# deep learning network trained in TensorFlow Object Detection API.
|
||||
# Then you can import it with a binary frozen graph (.pb) using readNetFromTensorflow() function.
|
||||
# See details and examples on the following wiki page: https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API
|
||||
import argparse
|
||||
import re
|
||||
from math import sqrt
|
||||
from tf_text_graph_common import *
|
||||
|
||||
class SSDAnchorGenerator:
|
||||
def __init__(self, min_scale, max_scale, num_layers, aspect_ratios,
|
||||
reduce_boxes_in_lowest_layer, image_width, image_height):
|
||||
self.min_scale = min_scale
|
||||
self.aspect_ratios = aspect_ratios
|
||||
self.reduce_boxes_in_lowest_layer = reduce_boxes_in_lowest_layer
|
||||
self.image_width = image_width
|
||||
self.image_height = image_height
|
||||
self.scales = [min_scale + (max_scale - min_scale) * i / (num_layers - 1)
|
||||
for i in range(num_layers)] + [1.0]
|
||||
|
||||
def get(self, layer_id):
|
||||
if layer_id == 0 and self.reduce_boxes_in_lowest_layer:
|
||||
widths = [0.1, self.min_scale * sqrt(2.0), self.min_scale * sqrt(0.5)]
|
||||
heights = [0.1, self.min_scale / sqrt(2.0), self.min_scale / sqrt(0.5)]
|
||||
else:
|
||||
widths = [self.scales[layer_id] * sqrt(ar) for ar in self.aspect_ratios]
|
||||
heights = [self.scales[layer_id] / sqrt(ar) for ar in self.aspect_ratios]
|
||||
|
||||
widths += [sqrt(self.scales[layer_id] * self.scales[layer_id + 1])]
|
||||
heights += [sqrt(self.scales[layer_id] * self.scales[layer_id + 1])]
|
||||
min_size = min(self.image_width, self.image_height)
|
||||
widths = [w * min_size for w in widths]
|
||||
heights = [h * min_size for h in heights]
|
||||
return widths, heights
|
||||
|
||||
|
||||
class MultiscaleAnchorGenerator:
|
||||
def __init__(self, min_level, aspect_ratios, scales_per_octave, anchor_scale):
|
||||
self.min_level = min_level
|
||||
self.aspect_ratios = aspect_ratios
|
||||
self.anchor_scale = anchor_scale
|
||||
self.scales = [2**(float(s) / scales_per_octave) for s in range(scales_per_octave)]
|
||||
|
||||
def get(self, layer_id):
|
||||
widths = []
|
||||
heights = []
|
||||
for a in self.aspect_ratios:
|
||||
for s in self.scales:
|
||||
base_anchor_size = 2**(self.min_level + layer_id) * self.anchor_scale
|
||||
ar = sqrt(a)
|
||||
heights.append(base_anchor_size * s / ar)
|
||||
widths.append(base_anchor_size * s * ar)
|
||||
return widths, heights
|
||||
|
||||
|
||||
def createSSDGraph(modelPath, configPath, outputPath):
|
||||
# Nodes that should be kept.
|
||||
keepOps = ['Conv2D', 'BiasAdd', 'Add', 'AddV2', 'Relu', 'Relu6', 'Placeholder', 'FusedBatchNorm',
|
||||
'DepthwiseConv2dNative', 'ConcatV2', 'Mul', 'MaxPool', 'AvgPool', 'Identity',
|
||||
'Sub', 'ResizeNearestNeighbor', 'Pad', 'FusedBatchNormV3', 'Mean']
|
||||
|
||||
# Node with which prefixes should be removed
|
||||
prefixesToRemove = ('MultipleGridAnchorGenerator/', 'Concatenate/', 'Postprocessor/', 'Preprocessor/map')
|
||||
|
||||
# Load a config file.
|
||||
config = readTextMessage(configPath)
|
||||
config = config['model'][0]['ssd'][0]
|
||||
num_classes = int(config['num_classes'][0])
|
||||
|
||||
fixed_shape_resizer = config['image_resizer'][0]['fixed_shape_resizer'][0]
|
||||
image_width = int(fixed_shape_resizer['width'][0])
|
||||
image_height = int(fixed_shape_resizer['height'][0])
|
||||
|
||||
box_predictor = 'convolutional' if 'convolutional_box_predictor' in config['box_predictor'][0] else 'weight_shared_convolutional'
|
||||
|
||||
anchor_generator = config['anchor_generator'][0]
|
||||
if 'ssd_anchor_generator' in anchor_generator:
|
||||
ssd_anchor_generator = anchor_generator['ssd_anchor_generator'][0]
|
||||
min_scale = float(ssd_anchor_generator['min_scale'][0])
|
||||
max_scale = float(ssd_anchor_generator['max_scale'][0])
|
||||
num_layers = int(ssd_anchor_generator['num_layers'][0])
|
||||
aspect_ratios = [float(ar) for ar in ssd_anchor_generator['aspect_ratios']]
|
||||
reduce_boxes_in_lowest_layer = True
|
||||
if 'reduce_boxes_in_lowest_layer' in ssd_anchor_generator:
|
||||
reduce_boxes_in_lowest_layer = ssd_anchor_generator['reduce_boxes_in_lowest_layer'][0] == 'true'
|
||||
priors_generator = SSDAnchorGenerator(min_scale, max_scale, num_layers,
|
||||
aspect_ratios, reduce_boxes_in_lowest_layer,
|
||||
image_width, image_height)
|
||||
|
||||
|
||||
print('Scale: [%f-%f]' % (min_scale, max_scale))
|
||||
print('Aspect ratios: %s' % str(aspect_ratios))
|
||||
print('Reduce boxes in the lowest layer: %s' % str(reduce_boxes_in_lowest_layer))
|
||||
elif 'multiscale_anchor_generator' in anchor_generator:
|
||||
multiscale_anchor_generator = anchor_generator['multiscale_anchor_generator'][0]
|
||||
min_level = int(multiscale_anchor_generator['min_level'][0])
|
||||
max_level = int(multiscale_anchor_generator['max_level'][0])
|
||||
anchor_scale = float(multiscale_anchor_generator['anchor_scale'][0])
|
||||
aspect_ratios = [float(ar) for ar in multiscale_anchor_generator['aspect_ratios']]
|
||||
scales_per_octave = int(multiscale_anchor_generator['scales_per_octave'][0])
|
||||
num_layers = max_level - min_level + 1
|
||||
priors_generator = MultiscaleAnchorGenerator(min_level, aspect_ratios,
|
||||
scales_per_octave, anchor_scale)
|
||||
print('Levels: [%d-%d]' % (min_level, max_level))
|
||||
print('Anchor scale: %f' % anchor_scale)
|
||||
print('Scales per octave: %d' % scales_per_octave)
|
||||
print('Aspect ratios: %s' % str(aspect_ratios))
|
||||
else:
|
||||
print('Unknown anchor_generator')
|
||||
exit(0)
|
||||
|
||||
print('Number of classes: %d' % num_classes)
|
||||
print('Number of layers: %d' % num_layers)
|
||||
print('box predictor: %s' % box_predictor)
|
||||
print('Input image size: %dx%d' % (image_width, image_height))
|
||||
|
||||
# Read the graph.
|
||||
outNames = ['num_detections', 'detection_scores', 'detection_boxes', 'detection_classes']
|
||||
|
||||
writeTextGraph(modelPath, outputPath, outNames)
|
||||
graph_def = parseTextGraph(outputPath)
|
||||
|
||||
def getUnconnectedNodes():
|
||||
unconnected = []
|
||||
for node in graph_def.node:
|
||||
unconnected.append(node.name)
|
||||
for inp in node.input:
|
||||
if inp in unconnected:
|
||||
unconnected.remove(inp)
|
||||
return unconnected
|
||||
|
||||
|
||||
def fuse_nodes(nodesToKeep):
|
||||
# Detect unfused batch normalization nodes and fuse them.
|
||||
# Add_0 <-- moving_variance, add_y
|
||||
# Rsqrt <-- Add_0
|
||||
# Mul_0 <-- Rsqrt, gamma
|
||||
# Mul_1 <-- input, Mul_0
|
||||
# Mul_2 <-- moving_mean, Mul_0
|
||||
# Sub_0 <-- beta, Mul_2
|
||||
# Add_1 <-- Mul_1, Sub_0
|
||||
nodesMap = {node.name: node for node in graph_def.node}
|
||||
subgraphBatchNorm = ['Add',
|
||||
['Mul', 'input', ['Mul', ['Rsqrt', ['Add', 'moving_variance', 'add_y']], 'gamma']],
|
||||
['Sub', 'beta', ['Mul', 'moving_mean', 'Mul_0']]]
|
||||
subgraphBatchNormV2 = ['AddV2',
|
||||
['Mul', 'input', ['Mul', ['Rsqrt', ['AddV2', 'moving_variance', 'add_y']], 'gamma']],
|
||||
['Sub', 'beta', ['Mul', 'moving_mean', 'Mul_0']]]
|
||||
# Detect unfused nearest neighbor resize.
|
||||
subgraphResizeNN = ['Reshape',
|
||||
['Mul', ['Reshape', 'input', ['Pack', 'shape_1', 'shape_2', 'shape_3', 'shape_4', 'shape_5']],
|
||||
'ones'],
|
||||
['Pack', ['StridedSlice', ['Shape', 'input'], 'stack', 'stack_1', 'stack_2'],
|
||||
'out_height', 'out_width', 'out_channels']]
|
||||
def checkSubgraph(node, targetNode, inputs, fusedNodes):
|
||||
op = targetNode[0]
|
||||
if node.op == op and (len(node.input) >= len(targetNode) - 1):
|
||||
fusedNodes.append(node)
|
||||
for i, inpOp in enumerate(targetNode[1:]):
|
||||
if isinstance(inpOp, list):
|
||||
if not node.input[i] in nodesMap or \
|
||||
not checkSubgraph(nodesMap[node.input[i]], inpOp, inputs, fusedNodes):
|
||||
return False
|
||||
else:
|
||||
inputs[inpOp] = node.input[i]
|
||||
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
nodesToRemove = []
|
||||
for node in graph_def.node:
|
||||
inputs = {}
|
||||
fusedNodes = []
|
||||
if checkSubgraph(node, subgraphBatchNorm, inputs, fusedNodes) or \
|
||||
checkSubgraph(node, subgraphBatchNormV2, inputs, fusedNodes):
|
||||
name = node.name
|
||||
node.Clear()
|
||||
node.name = name
|
||||
node.op = 'FusedBatchNorm'
|
||||
node.input.append(inputs['input'])
|
||||
node.input.append(inputs['gamma'])
|
||||
node.input.append(inputs['beta'])
|
||||
node.input.append(inputs['moving_mean'])
|
||||
node.input.append(inputs['moving_variance'])
|
||||
node.addAttr('epsilon', 0.001)
|
||||
nodesToRemove += fusedNodes[1:]
|
||||
|
||||
inputs = {}
|
||||
fusedNodes = []
|
||||
if checkSubgraph(node, subgraphResizeNN, inputs, fusedNodes):
|
||||
name = node.name
|
||||
node.Clear()
|
||||
node.name = name
|
||||
node.op = 'ResizeNearestNeighbor'
|
||||
node.input.append(inputs['input'])
|
||||
node.input.append(name + '/output_shape')
|
||||
|
||||
out_height_node = nodesMap[inputs['out_height']]
|
||||
out_width_node = nodesMap[inputs['out_width']]
|
||||
out_height = int(out_height_node.attr['value']['tensor'][0]['int_val'][0])
|
||||
out_width = int(out_width_node.attr['value']['tensor'][0]['int_val'][0])
|
||||
|
||||
shapeNode = NodeDef()
|
||||
shapeNode.name = name + '/output_shape'
|
||||
shapeNode.op = 'Const'
|
||||
shapeNode.addAttr('value', [out_height, out_width])
|
||||
graph_def.node.insert(graph_def.node.index(node), shapeNode)
|
||||
nodesToKeep.append(shapeNode.name)
|
||||
|
||||
nodesToRemove += fusedNodes[1:]
|
||||
for node in nodesToRemove:
|
||||
graph_def.node.remove(node)
|
||||
|
||||
nodesToKeep = []
|
||||
fuse_nodes(nodesToKeep)
|
||||
|
||||
removeIdentity(graph_def)
|
||||
|
||||
def to_remove(name, op):
|
||||
return (not name in nodesToKeep) and \
|
||||
(op == 'Const' or (not op in keepOps) or name.startswith(prefixesToRemove))
|
||||
|
||||
removeUnusedNodesAndAttrs(to_remove, graph_def)
|
||||
|
||||
|
||||
# Connect input node to the first layer
|
||||
assert(graph_def.node[0].op == 'Placeholder')
|
||||
try:
|
||||
input_shape = graph_def.node[0].attr['shape']['shape'][0]['dim']
|
||||
input_shape[1]['size'] = image_height
|
||||
input_shape[2]['size'] = image_width
|
||||
except:
|
||||
print("Input shapes are undefined")
|
||||
# assert(graph_def.node[1].op == 'Conv2D')
|
||||
weights = graph_def.node[1].input[-1]
|
||||
for i in range(len(graph_def.node[1].input)):
|
||||
graph_def.node[1].input.pop()
|
||||
graph_def.node[1].input.append(graph_def.node[0].name)
|
||||
graph_def.node[1].input.append(weights)
|
||||
|
||||
# check and correct the case when preprocessing block is after input
|
||||
preproc_id = "Preprocessor/"
|
||||
if graph_def.node[2].name.startswith(preproc_id) and \
|
||||
graph_def.node[2].input[0].startswith(preproc_id):
|
||||
|
||||
if not any(preproc_id in inp for inp in graph_def.node[3].input):
|
||||
graph_def.node[3].input.insert(0, graph_def.node[2].name)
|
||||
|
||||
|
||||
# Create SSD postprocessing head ###############################################
|
||||
|
||||
# Concatenate predictions of classes, predictions of bounding boxes and proposals.
|
||||
def addConcatNode(name, inputs, axisNodeName):
|
||||
concat = NodeDef()
|
||||
concat.name = name
|
||||
concat.op = 'ConcatV2'
|
||||
for inp in inputs:
|
||||
concat.input.append(inp)
|
||||
concat.input.append(axisNodeName)
|
||||
graph_def.node.extend([concat])
|
||||
|
||||
addConstNode('concat/axis_flatten', [-1], graph_def)
|
||||
addConstNode('PriorBox/concat/axis', [-2], graph_def)
|
||||
|
||||
for label in ['ClassPredictor', 'BoxEncodingPredictor' if box_predictor == 'convolutional' else 'BoxPredictor']:
|
||||
concatInputs = []
|
||||
for i in range(num_layers):
|
||||
# Flatten predictions
|
||||
flatten = NodeDef()
|
||||
if box_predictor == 'convolutional':
|
||||
inpName = 'BoxPredictor_%d/%s/BiasAdd' % (i, label)
|
||||
else:
|
||||
if i == 0:
|
||||
inpName = 'WeightSharedConvolutionalBoxPredictor/%s/BiasAdd' % label
|
||||
else:
|
||||
inpName = 'WeightSharedConvolutionalBoxPredictor_%d/%s/BiasAdd' % (i, label)
|
||||
flatten.input.append(inpName)
|
||||
flatten.name = inpName + '/Flatten'
|
||||
flatten.op = 'Flatten'
|
||||
|
||||
concatInputs.append(flatten.name)
|
||||
graph_def.node.extend([flatten])
|
||||
addConcatNode('%s/concat' % label, concatInputs, 'concat/axis_flatten')
|
||||
|
||||
num_matched_layers = 0
|
||||
for node in graph_def.node:
|
||||
if re.match('BoxPredictor_\d/BoxEncodingPredictor/convolution', node.name) or \
|
||||
re.match('BoxPredictor_\d/BoxEncodingPredictor/Conv2D', node.name) or \
|
||||
re.match('WeightSharedConvolutionalBoxPredictor(_\d)*/BoxPredictor/Conv2D', node.name):
|
||||
node.addAttr('loc_pred_transposed', True)
|
||||
num_matched_layers += 1
|
||||
assert(num_matched_layers == num_layers)
|
||||
|
||||
# Add layers that generate anchors (bounding boxes proposals).
|
||||
priorBoxes = []
|
||||
boxCoder = config['box_coder'][0]
|
||||
fasterRcnnBoxCoder = boxCoder['faster_rcnn_box_coder'][0]
|
||||
boxCoderVariance = [1.0/float(fasterRcnnBoxCoder['x_scale'][0]), 1.0/float(fasterRcnnBoxCoder['y_scale'][0]), 1.0/float(fasterRcnnBoxCoder['width_scale'][0]), 1.0/float(fasterRcnnBoxCoder['height_scale'][0])]
|
||||
for i in range(num_layers):
|
||||
priorBox = NodeDef()
|
||||
priorBox.name = 'PriorBox_%d' % i
|
||||
priorBox.op = 'PriorBox'
|
||||
if box_predictor == 'convolutional':
|
||||
priorBox.input.append('BoxPredictor_%d/BoxEncodingPredictor/BiasAdd' % i)
|
||||
else:
|
||||
if i == 0:
|
||||
priorBox.input.append('WeightSharedConvolutionalBoxPredictor/BoxPredictor/Conv2D')
|
||||
else:
|
||||
priorBox.input.append('WeightSharedConvolutionalBoxPredictor_%d/BoxPredictor/BiasAdd' % i)
|
||||
priorBox.input.append(graph_def.node[0].name) # image_tensor
|
||||
|
||||
priorBox.addAttr('flip', False)
|
||||
priorBox.addAttr('clip', False)
|
||||
|
||||
widths, heights = priors_generator.get(i)
|
||||
|
||||
priorBox.addAttr('width', widths)
|
||||
priorBox.addAttr('height', heights)
|
||||
priorBox.addAttr('variance', boxCoderVariance)
|
||||
|
||||
graph_def.node.extend([priorBox])
|
||||
priorBoxes.append(priorBox.name)
|
||||
|
||||
# Compare this layer's output with Postprocessor/Reshape
|
||||
addConcatNode('PriorBox/concat', priorBoxes, 'concat/axis_flatten')
|
||||
|
||||
# Sigmoid for classes predictions and DetectionOutput layer
|
||||
addReshape('ClassPredictor/concat', 'ClassPredictor/concat3d', [0, -1, num_classes + 1], graph_def)
|
||||
|
||||
sigmoid = NodeDef()
|
||||
sigmoid.name = 'ClassPredictor/concat/sigmoid'
|
||||
sigmoid.op = 'Sigmoid'
|
||||
sigmoid.input.append('ClassPredictor/concat3d')
|
||||
graph_def.node.extend([sigmoid])
|
||||
|
||||
addFlatten(sigmoid.name, sigmoid.name + '/Flatten', graph_def)
|
||||
|
||||
detectionOut = NodeDef()
|
||||
detectionOut.name = 'detection_out'
|
||||
detectionOut.op = 'DetectionOutput'
|
||||
|
||||
if box_predictor == 'convolutional':
|
||||
detectionOut.input.append('BoxEncodingPredictor/concat')
|
||||
else:
|
||||
detectionOut.input.append('BoxPredictor/concat')
|
||||
detectionOut.input.append(sigmoid.name + '/Flatten')
|
||||
detectionOut.input.append('PriorBox/concat')
|
||||
|
||||
detectionOut.addAttr('num_classes', num_classes + 1)
|
||||
detectionOut.addAttr('share_location', True)
|
||||
detectionOut.addAttr('background_label_id', 0)
|
||||
|
||||
postProcessing = config['post_processing'][0]
|
||||
batchNMS = postProcessing['batch_non_max_suppression'][0]
|
||||
|
||||
if 'iou_threshold' in batchNMS:
|
||||
detectionOut.addAttr('nms_threshold', float(batchNMS['iou_threshold'][0]))
|
||||
else:
|
||||
detectionOut.addAttr('nms_threshold', 0.6)
|
||||
|
||||
if 'score_threshold' in batchNMS:
|
||||
detectionOut.addAttr('confidence_threshold', float(batchNMS['score_threshold'][0]))
|
||||
else:
|
||||
detectionOut.addAttr('confidence_threshold', 0.01)
|
||||
|
||||
if 'max_detections_per_class' in batchNMS:
|
||||
detectionOut.addAttr('top_k', int(batchNMS['max_detections_per_class'][0]))
|
||||
else:
|
||||
detectionOut.addAttr('top_k', 100)
|
||||
|
||||
if 'max_total_detections' in batchNMS:
|
||||
detectionOut.addAttr('keep_top_k', int(batchNMS['max_total_detections'][0]))
|
||||
else:
|
||||
detectionOut.addAttr('keep_top_k', 100)
|
||||
|
||||
detectionOut.addAttr('code_type', "CENTER_SIZE")
|
||||
|
||||
graph_def.node.extend([detectionOut])
|
||||
|
||||
while True:
|
||||
unconnectedNodes = getUnconnectedNodes()
|
||||
unconnectedNodes.remove(detectionOut.name)
|
||||
if not unconnectedNodes:
|
||||
break
|
||||
|
||||
for name in unconnectedNodes:
|
||||
for i in range(len(graph_def.node)):
|
||||
if graph_def.node[i].name == name:
|
||||
del graph_def.node[i]
|
||||
break
|
||||
|
||||
# Save as text.
|
||||
graph_def.save(outputPath)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Run this script to get a text graph of '
|
||||
'SSD model from TensorFlow Object Detection API. '
|
||||
'Then pass it with .pb file to cv::dnn::readNetFromTensorflow function.')
|
||||
parser.add_argument('--input', required=True, help='Path to frozen TensorFlow graph.')
|
||||
parser.add_argument('--output', required=True, help='Path to output text graph.')
|
||||
parser.add_argument('--config', required=True, help='Path to a *.config file is used for training.')
|
||||
args = parser.parse_args()
|
||||
|
||||
createSSDGraph(args.input, args.config, args.output)
|
||||
@@ -0,0 +1,471 @@
|
||||
#!/usr/bin/env python3
|
||||
'''
|
||||
You can download the Geometric Matching Module model from https://www.dropbox.com/s/tyhc73xa051grjp/cp_vton_gmm.onnx?dl=0
|
||||
You can download the Try-On Module model from https://www.dropbox.com/s/q2x97ve2h53j66k/cp_vton_tom.onnx?dl=0
|
||||
You can download the cloth segmentation model from https://www.dropbox.com/s/qag9vzambhhkvxr/lip_jppnet_384.pb?dl=0
|
||||
You can find the OpenPose proto in opencv_extra/testdata/dnn/openpose_pose_coco.prototxt
|
||||
and get .caffemodel using opencv_extra/testdata/dnn/download_models.py
|
||||
'''
|
||||
|
||||
import argparse
|
||||
import os.path
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
|
||||
from numpy import linalg
|
||||
from common import findFile
|
||||
from human_parsing import parse_human
|
||||
|
||||
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV,
|
||||
cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
|
||||
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD, cv.dnn.DNN_TARGET_HDDL,
|
||||
cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Use this script to run virtial try-on using CP-VTON',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input_image', '-i', required=True, help='Path to image with person.')
|
||||
parser.add_argument('--input_cloth', '-c', required=True, help='Path to target cloth image')
|
||||
parser.add_argument('--gmm_model', '-gmm', default='cp_vton_gmm.onnx', help='Path to Geometric Matching Module .onnx model.')
|
||||
parser.add_argument('--tom_model', '-tom', default='cp_vton_tom.onnx', help='Path to Try-On Module .onnx model.')
|
||||
parser.add_argument('--segmentation_model', default='lip_jppnet_384.pb', help='Path to cloth segmentation .pb model.')
|
||||
parser.add_argument('--openpose_proto', default='openpose_pose_coco.prototxt', help='Path to OpenPose .prototxt model was trained on COCO dataset.')
|
||||
parser.add_argument('--openpose_model', default='openpose_pose_coco.caffemodel', help='Path to OpenPose .caffemodel model was trained on COCO dataset.')
|
||||
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
|
||||
help="Choose one of computation backends: "
|
||||
"%d: automatically (by default), "
|
||||
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"%d: OpenCV implementation, "
|
||||
"%d: VKCOM, "
|
||||
"%d: CUDA" % backends)
|
||||
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
|
||||
help='Choose one of target computation devices: '
|
||||
'%d: CPU target (by default), '
|
||||
'%d: OpenCL, '
|
||||
'%d: OpenCL fp16 (half-float precision), '
|
||||
'%d: NCS2 VPU, '
|
||||
'%d: HDDL VPU, '
|
||||
'%d: Vulkan, '
|
||||
'%d: CUDA, '
|
||||
'%d: CUDA fp16 (half-float preprocess)'% targets)
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
|
||||
def get_pose_map(image, proto_path, model_path, backend, target, height=256, width=192):
|
||||
radius = 5
|
||||
inp = cv.dnn.blobFromImage(image, 1.0 / 255, (width, height))
|
||||
|
||||
net = cv.dnn.readNet(proto_path, model_path)
|
||||
net.setPreferableBackend(backend)
|
||||
net.setPreferableTarget(target)
|
||||
net.setInput(inp)
|
||||
out = net.forward()
|
||||
|
||||
threshold = 0.1
|
||||
_, out_c, out_h, out_w = out.shape
|
||||
pose_map = np.zeros((height, width, out_c - 1))
|
||||
# last label: Background
|
||||
for i in range(0, out.shape[1] - 1):
|
||||
heatMap = out[0, i, :, :]
|
||||
keypoint = np.full((height, width), -1)
|
||||
_, conf, _, point = cv.minMaxLoc(heatMap)
|
||||
x = width * point[0] // out_w
|
||||
y = height * point[1] // out_h
|
||||
if conf > threshold and x > 0 and y > 0:
|
||||
keypoint[y - radius:y + radius, x - radius:x + radius] = 1
|
||||
pose_map[:, :, i] = keypoint
|
||||
|
||||
pose_map = pose_map.transpose(2, 0, 1)
|
||||
return pose_map
|
||||
|
||||
|
||||
class BilinearFilter(object):
|
||||
"""
|
||||
PIL bilinear resize implementation
|
||||
image = image.resize((image_width // 16, image_height // 16), Image.BILINEAR)
|
||||
"""
|
||||
def _precompute_coeffs(self, inSize, outSize):
|
||||
filterscale = max(1.0, inSize / outSize)
|
||||
ksize = int(np.ceil(filterscale)) * 2 + 1
|
||||
|
||||
kk = np.zeros(shape=(outSize * ksize, ), dtype=np.float32)
|
||||
bounds = np.empty(shape=(outSize * 2, ), dtype=np.int32)
|
||||
|
||||
centers = (np.arange(outSize) + 0.5) * filterscale + 0.5
|
||||
bounds[::2] = np.where(centers - filterscale < 0, 0, centers - filterscale)
|
||||
bounds[1::2] = np.where(centers + filterscale > inSize, inSize, centers + filterscale) - bounds[::2]
|
||||
xmins = bounds[::2] - centers + 1
|
||||
|
||||
points = np.array([np.arange(row) + xmins[i] for i, row in enumerate(bounds[1::2])]) / filterscale
|
||||
for xx in range(0, outSize):
|
||||
point = points[xx]
|
||||
bilinear = np.where(point < 1.0, 1.0 - abs(point), 0.0)
|
||||
ww = np.sum(bilinear)
|
||||
kk[xx * ksize : xx * ksize + bilinear.size] = np.where(ww == 0.0, bilinear, bilinear / ww)
|
||||
return bounds, kk, ksize
|
||||
|
||||
def _resample_horizontal(self, out, img, ksize, bounds, kk):
|
||||
for yy in range(0, out.shape[0]):
|
||||
for xx in range(0, out.shape[1]):
|
||||
xmin = bounds[xx * 2 + 0]
|
||||
xmax = bounds[xx * 2 + 1]
|
||||
k = kk[xx * ksize : xx * ksize + xmax]
|
||||
out[yy, xx] = np.round(np.sum(img[yy, xmin : xmin + xmax] * k))
|
||||
|
||||
def _resample_vertical(self, out, img, ksize, bounds, kk):
|
||||
for yy in range(0, out.shape[0]):
|
||||
ymin = bounds[yy * 2 + 0]
|
||||
ymax = bounds[yy * 2 + 1]
|
||||
k = kk[yy * ksize: yy * ksize + ymax]
|
||||
out[yy] = np.round(np.sum(img[ymin : ymin + ymax, 0:out.shape[1]] * k[:, np.newaxis], axis=0))
|
||||
|
||||
def imaging_resample(self, img, xsize, ysize):
|
||||
height, width = img.shape[0:2]
|
||||
bounds_horiz, kk_horiz, ksize_horiz = self._precompute_coeffs(width, xsize)
|
||||
bounds_vert, kk_vert, ksize_vert = self._precompute_coeffs(height, ysize)
|
||||
|
||||
out_hor = np.empty((img.shape[0], xsize), dtype=np.uint8)
|
||||
self._resample_horizontal(out_hor, img, ksize_horiz, bounds_horiz, kk_horiz)
|
||||
out = np.empty((ysize, xsize), dtype=np.uint8)
|
||||
self._resample_vertical(out, out_hor, ksize_vert, bounds_vert, kk_vert)
|
||||
return out
|
||||
|
||||
|
||||
class CpVton(object):
|
||||
def __init__(self, gmm_model, tom_model, backend, target):
|
||||
super(CpVton, self).__init__()
|
||||
self.gmm_net = cv.dnn.readNet(gmm_model)
|
||||
self.tom_net = cv.dnn.readNet(tom_model)
|
||||
self.gmm_net.setPreferableBackend(backend)
|
||||
self.gmm_net.setPreferableTarget(target)
|
||||
self.tom_net.setPreferableBackend(backend)
|
||||
self.tom_net.setPreferableTarget(target)
|
||||
|
||||
def prepare_agnostic(self, segm_image, input_image, pose_map, height=256, width=192):
|
||||
palette = {
|
||||
'Background' : (0, 0, 0),
|
||||
'Hat' : (128, 0, 0),
|
||||
'Hair' : (255, 0, 0),
|
||||
'Glove' : (0, 85, 0),
|
||||
'Sunglasses' : (170, 0, 51),
|
||||
'UpperClothes' : (255, 85, 0),
|
||||
'Dress' : (0, 0, 85),
|
||||
'Coat' : (0, 119, 221),
|
||||
'Socks' : (85, 85, 0),
|
||||
'Pants' : (0, 85, 85),
|
||||
'Jumpsuits' : (85, 51, 0),
|
||||
'Scarf' : (52, 86, 128),
|
||||
'Skirt' : (0, 128, 0),
|
||||
'Face' : (0, 0, 255),
|
||||
'Left-arm' : (51, 170, 221),
|
||||
'Right-arm' : (0, 255, 255),
|
||||
'Left-leg' : (85, 255, 170),
|
||||
'Right-leg' : (170, 255, 85),
|
||||
'Left-shoe' : (255, 255, 0),
|
||||
'Right-shoe' : (255, 170, 0)
|
||||
}
|
||||
color2label = {val: key for key, val in palette.items()}
|
||||
head_labels = ['Hat', 'Hair', 'Sunglasses', 'Face', 'Pants', 'Skirt']
|
||||
|
||||
segm_image = cv.cvtColor(segm_image, cv.COLOR_BGR2RGB)
|
||||
phead = np.zeros((1, height, width), dtype=np.float32)
|
||||
pose_shape = np.zeros((height, width), dtype=np.uint8)
|
||||
for r in range(height):
|
||||
for c in range(width):
|
||||
pixel = tuple(segm_image[r, c])
|
||||
if tuple(pixel) in color2label:
|
||||
if color2label[pixel] in head_labels:
|
||||
phead[0, r, c] = 1
|
||||
if color2label[pixel] != 'Background':
|
||||
pose_shape[r, c] = 255
|
||||
|
||||
input_image = cv.dnn.blobFromImage(input_image, 1.0 / 127.5, (width, height), mean=(127.5, 127.5, 127.5), swapRB=True)
|
||||
input_image = input_image.squeeze(0)
|
||||
|
||||
img_head = input_image * phead - (1 - phead)
|
||||
|
||||
downsample = BilinearFilter()
|
||||
down = downsample.imaging_resample(pose_shape, width // 16, height // 16)
|
||||
res_shape = cv.resize(down, (width, height), cv.INTER_LINEAR)
|
||||
|
||||
res_shape = cv.dnn.blobFromImage(res_shape, 1.0 / 127.5, mean=(127.5, 127.5, 127.5), swapRB=True)
|
||||
res_shape = res_shape.squeeze(0)
|
||||
|
||||
agnostic = np.concatenate((res_shape, img_head, pose_map), axis=0)
|
||||
agnostic = np.expand_dims(agnostic, axis=0)
|
||||
return agnostic.astype(np.float32)
|
||||
|
||||
def get_warped_cloth(self, cloth_img, agnostic, height=256, width=192):
|
||||
cloth = cv.dnn.blobFromImage(cloth_img, 1.0 / 127.5, (width, height), mean=(127.5, 127.5, 127.5), swapRB=True)
|
||||
|
||||
self.gmm_net.setInput(agnostic, "input.1")
|
||||
self.gmm_net.setInput(cloth, "input.18")
|
||||
theta = self.gmm_net.forward()
|
||||
|
||||
grid = self._generate_grid(theta)
|
||||
warped_cloth = self._bilinear_sampler(cloth, grid).astype(np.float32)
|
||||
return warped_cloth
|
||||
|
||||
def get_tryon(self, agnostic, warp_cloth):
|
||||
inp = np.concatenate([agnostic, warp_cloth], axis=1)
|
||||
self.tom_net.setInput(inp)
|
||||
out = self.tom_net.forward()
|
||||
|
||||
p_rendered, m_composite = np.split(out, [3], axis=1)
|
||||
p_rendered = np.tanh(p_rendered)
|
||||
m_composite = 1 / (1 + np.exp(-m_composite))
|
||||
|
||||
p_tryon = warp_cloth * m_composite + p_rendered * (1 - m_composite)
|
||||
rgb_p_tryon = cv.cvtColor(p_tryon.squeeze(0).transpose(1, 2, 0), cv.COLOR_BGR2RGB)
|
||||
rgb_p_tryon = (rgb_p_tryon + 1) / 2
|
||||
return rgb_p_tryon
|
||||
|
||||
def _compute_L_inverse(self, X, Y):
|
||||
N = X.shape[0]
|
||||
|
||||
Xmat = np.tile(X, (1, N))
|
||||
Ymat = np.tile(Y, (1, N))
|
||||
P_dist_squared = np.power(Xmat - Xmat.transpose(1, 0), 2) + np.power(Ymat - Ymat.transpose(1, 0), 2)
|
||||
|
||||
P_dist_squared[P_dist_squared == 0] = 1
|
||||
K = np.multiply(P_dist_squared, np.log(P_dist_squared))
|
||||
|
||||
O = np.ones([N, 1], dtype=np.float32)
|
||||
Z = np.zeros([3, 3], dtype=np.float32)
|
||||
P = np.concatenate([O, X, Y], axis=1)
|
||||
first = np.concatenate((K, P), axis=1)
|
||||
second = np.concatenate((P.transpose(1, 0), Z), axis=1)
|
||||
L = np.concatenate((first, second), axis=0)
|
||||
Li = linalg.inv(L)
|
||||
return Li
|
||||
|
||||
def _prepare_to_transform(self, out_h=256, out_w=192, grid_size=5):
|
||||
grid_X, grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h))
|
||||
grid_X = np.expand_dims(np.expand_dims(grid_X, axis=0), axis=3)
|
||||
grid_Y = np.expand_dims(np.expand_dims(grid_Y, axis=0), axis=3)
|
||||
|
||||
axis_coords = np.linspace(-1, 1, grid_size)
|
||||
N = grid_size ** 2
|
||||
P_Y, P_X = np.meshgrid(axis_coords, axis_coords)
|
||||
|
||||
P_X = np.reshape(P_X,(-1, 1))
|
||||
P_Y = np.reshape(P_Y,(-1, 1))
|
||||
|
||||
P_X = np.expand_dims(np.expand_dims(np.expand_dims(P_X, axis=2), axis=3), axis=4).transpose(4, 1, 2, 3, 0)
|
||||
P_Y = np.expand_dims(np.expand_dims(np.expand_dims(P_Y, axis=2), axis=3), axis=4).transpose(4, 1, 2, 3, 0)
|
||||
return grid_X, grid_Y, N, P_X, P_Y
|
||||
|
||||
def _expand_torch(self, X, shape):
|
||||
if len(X.shape) != len(shape):
|
||||
return X.flatten().reshape(shape)
|
||||
else:
|
||||
axis = [1 if src == dst else dst for src, dst in zip(X.shape, shape)]
|
||||
return np.tile(X, axis)
|
||||
|
||||
def _apply_transformation(self, theta, points, N, P_X, P_Y):
|
||||
if len(theta.shape) == 2:
|
||||
theta = np.expand_dims(np.expand_dims(theta, axis=2), axis=3)
|
||||
|
||||
batch_size = theta.shape[0]
|
||||
|
||||
P_X_base = np.copy(P_X)
|
||||
P_Y_base = np.copy(P_Y)
|
||||
|
||||
Li = self._compute_L_inverse(np.reshape(P_X, (N, -1)), np.reshape(P_Y, (N, -1)))
|
||||
Li = np.expand_dims(Li, axis=0)
|
||||
|
||||
# split theta into point coordinates
|
||||
Q_X = np.squeeze(theta[:, :N, :, :], axis=3)
|
||||
Q_Y = np.squeeze(theta[:, N:, :, :], axis=3)
|
||||
|
||||
Q_X += self._expand_torch(P_X_base, Q_X.shape)
|
||||
Q_Y += self._expand_torch(P_Y_base, Q_Y.shape)
|
||||
|
||||
points_b = points.shape[0]
|
||||
points_h = points.shape[1]
|
||||
points_w = points.shape[2]
|
||||
|
||||
P_X = self._expand_torch(P_X, (1, points_h, points_w, 1, N))
|
||||
P_Y = self._expand_torch(P_Y, (1, points_h, points_w, 1, N))
|
||||
|
||||
W_X = self._expand_torch(Li[:,:N,:N], (batch_size, N, N)) @ Q_X
|
||||
W_Y = self._expand_torch(Li[:,:N,:N], (batch_size, N, N)) @ Q_Y
|
||||
|
||||
W_X = np.expand_dims(np.expand_dims(W_X, axis=3), axis=4).transpose(0, 4, 2, 3, 1)
|
||||
W_X = np.repeat(W_X, points_h, axis=1)
|
||||
W_X = np.repeat(W_X, points_w, axis=2)
|
||||
|
||||
W_Y = np.expand_dims(np.expand_dims(W_Y, axis=3), axis=4).transpose(0, 4, 2, 3, 1)
|
||||
W_Y = np.repeat(W_Y, points_h, axis=1)
|
||||
W_Y = np.repeat(W_Y, points_w, axis=2)
|
||||
|
||||
A_X = self._expand_torch(Li[:, N:, :N], (batch_size, 3, N)) @ Q_X
|
||||
A_Y = self._expand_torch(Li[:, N:, :N], (batch_size, 3, N)) @ Q_Y
|
||||
|
||||
A_X = np.expand_dims(np.expand_dims(A_X, axis=3), axis=4).transpose(0, 4, 2, 3, 1)
|
||||
A_X = np.repeat(A_X, points_h, axis=1)
|
||||
A_X = np.repeat(A_X, points_w, axis=2)
|
||||
|
||||
A_Y = np.expand_dims(np.expand_dims(A_Y, axis=3), axis=4).transpose(0, 4, 2, 3, 1)
|
||||
A_Y = np.repeat(A_Y, points_h, axis=1)
|
||||
A_Y = np.repeat(A_Y, points_w, axis=2)
|
||||
|
||||
points_X_for_summation = np.expand_dims(np.expand_dims(points[:, :, :, 0], axis=3), axis=4)
|
||||
points_X_for_summation = self._expand_torch(points_X_for_summation, points[:, :, :, 0].shape + (1, N))
|
||||
|
||||
points_Y_for_summation = np.expand_dims(np.expand_dims(points[:, :, :, 1], axis=3), axis=4)
|
||||
points_Y_for_summation = self._expand_torch(points_Y_for_summation, points[:, :, :, 0].shape + (1, N))
|
||||
|
||||
if points_b == 1:
|
||||
delta_X = points_X_for_summation - P_X
|
||||
delta_Y = points_Y_for_summation - P_Y
|
||||
else:
|
||||
delta_X = points_X_for_summation - self._expand_torch(P_X, points_X_for_summation.shape)
|
||||
delta_Y = points_Y_for_summation - self._expand_torch(P_Y, points_Y_for_summation.shape)
|
||||
|
||||
dist_squared = np.power(delta_X, 2) + np.power(delta_Y, 2)
|
||||
dist_squared[dist_squared == 0] = 1
|
||||
U = np.multiply(dist_squared, np.log(dist_squared))
|
||||
|
||||
points_X_batch = np.expand_dims(points[:,:,:,0], axis=3)
|
||||
points_Y_batch = np.expand_dims(points[:,:,:,1], axis=3)
|
||||
|
||||
if points_b == 1:
|
||||
points_X_batch = self._expand_torch(points_X_batch, (batch_size, ) + points_X_batch.shape[1:])
|
||||
points_Y_batch = self._expand_torch(points_Y_batch, (batch_size, ) + points_Y_batch.shape[1:])
|
||||
|
||||
points_X_prime = A_X[:,:,:,:,0]+ \
|
||||
np.multiply(A_X[:,:,:,:,1], points_X_batch) + \
|
||||
np.multiply(A_X[:,:,:,:,2], points_Y_batch) + \
|
||||
np.sum(np.multiply(W_X, self._expand_torch(U, W_X.shape)), 4)
|
||||
|
||||
points_Y_prime = A_Y[:,:,:,:,0]+ \
|
||||
np.multiply(A_Y[:,:,:,:,1], points_X_batch) + \
|
||||
np.multiply(A_Y[:,:,:,:,2], points_Y_batch) + \
|
||||
np.sum(np.multiply(W_Y, self._expand_torch(U, W_Y.shape)), 4)
|
||||
|
||||
return np.concatenate((points_X_prime, points_Y_prime), 3)
|
||||
|
||||
def _generate_grid(self, theta):
|
||||
grid_X, grid_Y, N, P_X, P_Y = self._prepare_to_transform()
|
||||
warped_grid = self._apply_transformation(theta, np.concatenate((grid_X, grid_Y), axis=3), N, P_X, P_Y)
|
||||
return warped_grid
|
||||
|
||||
def _bilinear_sampler(self, img, grid):
|
||||
x, y = grid[:,:,:,0], grid[:,:,:,1]
|
||||
|
||||
H = img.shape[2]
|
||||
W = img.shape[3]
|
||||
max_y = H - 1
|
||||
max_x = W - 1
|
||||
|
||||
# rescale x and y to [0, W-1/H-1]
|
||||
x = 0.5 * (x + 1.0) * (max_x - 1)
|
||||
y = 0.5 * (y + 1.0) * (max_y - 1)
|
||||
|
||||
# grab 4 nearest corner points for each (x_i, y_i)
|
||||
x0 = np.floor(x).astype(int)
|
||||
x1 = x0 + 1
|
||||
y0 = np.floor(y).astype(int)
|
||||
y1 = y0 + 1
|
||||
|
||||
# calculate deltas
|
||||
wa = (x1 - x) * (y1 - y)
|
||||
wb = (x1 - x) * (y - y0)
|
||||
wc = (x - x0) * (y1 - y)
|
||||
wd = (x - x0) * (y - y0)
|
||||
|
||||
# clip to range [0, H-1/W-1] to not violate img boundaries
|
||||
x0 = np.clip(x0, 0, max_x)
|
||||
x1 = np.clip(x1, 0, max_x)
|
||||
y0 = np.clip(y0, 0, max_y)
|
||||
y1 = np.clip(y1, 0, max_y)
|
||||
|
||||
# get pixel value at corner coords
|
||||
img = img.reshape(-1, H, W)
|
||||
Ia = img[:, y0, x0].swapaxes(0, 1)
|
||||
Ib = img[:, y1, x0].swapaxes(0, 1)
|
||||
Ic = img[:, y0, x1].swapaxes(0, 1)
|
||||
Id = img[:, y1, x1].swapaxes(0, 1)
|
||||
|
||||
wa = np.expand_dims(wa, axis=0)
|
||||
wb = np.expand_dims(wb, axis=0)
|
||||
wc = np.expand_dims(wc, axis=0)
|
||||
wd = np.expand_dims(wd, axis=0)
|
||||
|
||||
# compute output
|
||||
out = wa*Ia + wb*Ib + wc*Ic + wd*Id
|
||||
return out
|
||||
|
||||
|
||||
class CorrelationLayer(object):
|
||||
def __init__(self, params, blobs):
|
||||
super(CorrelationLayer, self).__init__()
|
||||
|
||||
def getMemoryShapes(self, inputs):
|
||||
fetureAShape = inputs[0]
|
||||
b, _, h, w = fetureAShape
|
||||
return [[b, h * w, h, w]]
|
||||
|
||||
def forward(self, inputs):
|
||||
feature_A, feature_B = inputs
|
||||
b, c, h, w = feature_A.shape
|
||||
feature_A = feature_A.transpose(0, 1, 3, 2)
|
||||
feature_A = np.reshape(feature_A, (b, c, h * w))
|
||||
feature_B = np.reshape(feature_B, (b, c, h * w))
|
||||
feature_B = feature_B.transpose(0, 2, 1)
|
||||
feature_mul = feature_B @ feature_A
|
||||
feature_mul= np.reshape(feature_mul, (b, h, w, h * w))
|
||||
feature_mul = feature_mul.transpose(0, 1, 3, 2)
|
||||
correlation_tensor = feature_mul.transpose(0, 2, 1, 3)
|
||||
correlation_tensor = np.ascontiguousarray(correlation_tensor)
|
||||
return [correlation_tensor]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not os.path.isfile(args.gmm_model):
|
||||
raise OSError("GMM model not exist")
|
||||
if not os.path.isfile(args.tom_model):
|
||||
raise OSError("TOM model not exist")
|
||||
if not os.path.isfile(args.segmentation_model):
|
||||
raise OSError("Segmentation model not exist")
|
||||
if not os.path.isfile(findFile(args.openpose_proto)):
|
||||
raise OSError("OpenPose proto not exist")
|
||||
if not os.path.isfile(findFile(args.openpose_model)):
|
||||
raise OSError("OpenPose model not exist")
|
||||
|
||||
person_img = cv.imread(args.input_image)
|
||||
ratio = 256 / 192
|
||||
inp_h, inp_w, _ = person_img.shape
|
||||
current_ratio = inp_h / inp_w
|
||||
if current_ratio > ratio:
|
||||
center_h = inp_h // 2
|
||||
out_h = inp_w * ratio
|
||||
start = int(center_h - out_h // 2)
|
||||
end = int(center_h + out_h // 2)
|
||||
person_img = person_img[start:end, ...]
|
||||
else:
|
||||
center_w = inp_w // 2
|
||||
out_w = inp_h / ratio
|
||||
start = int(center_w - out_w // 2)
|
||||
end = int(center_w + out_w // 2)
|
||||
person_img = person_img[:, start:end, :]
|
||||
|
||||
cloth_img = cv.imread(args.input_cloth)
|
||||
pose = get_pose_map(person_img, findFile(args.openpose_proto),
|
||||
findFile(args.openpose_model), args.backend, args.target)
|
||||
segm_image = parse_human(person_img, args.segmentation_model)
|
||||
segm_image = cv.resize(segm_image, (192, 256), cv.INTER_LINEAR)
|
||||
|
||||
cv.dnn_registerLayer('Correlation', CorrelationLayer)
|
||||
|
||||
model = CpVton(args.gmm_model, args.tom_model, args.backend, args.target)
|
||||
agnostic = model.prepare_agnostic(segm_image, person_img, pose)
|
||||
warped_cloth = model.get_warped_cloth(cloth_img, agnostic)
|
||||
output = model.get_tryon(agnostic, warped_cloth)
|
||||
|
||||
cv.dnn_unregisterLayer('Correlation')
|
||||
|
||||
winName = 'Virtual Try-On'
|
||||
cv.namedWindow(winName, cv.WINDOW_AUTOSIZE)
|
||||
cv.imshow(winName, output)
|
||||
cv.waitKey()
|
||||
@@ -0,0 +1,130 @@
|
||||
# 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) 2026, BigVision LLC, all rights reserved.
|
||||
# Third party copyrights are property of their respective owners.
|
||||
|
||||
'''
|
||||
This is a sample script to run PaliGemma2 vision-language inference in OpenCV using
|
||||
ONNX models. Given an image and a text prompt, it generates a text response
|
||||
(e.g. a caption).
|
||||
|
||||
The model is split into three ONNX files:
|
||||
- SigLIP vision encoder : image -> 256 image-feature tokens
|
||||
- Embedding : prompt token ids -> text embeddings
|
||||
- Gemma2 language model : [image_features | text_embeds] -> logits
|
||||
|
||||
Model: https://huggingface.co/google/paligemma2-3b-pt-224
|
||||
ONNX: https://huggingface.co/nklskyoy/paligemma2-3b-pt-224-onnx
|
||||
|
||||
Run the script:
|
||||
1. Install the required dependencies:
|
||||
|
||||
pip install numpy
|
||||
|
||||
2. Run the script:
|
||||
|
||||
python vlm_inference.py --siglip=<path-to-vision_model.onnx> \
|
||||
--embedding=<path-to-embedding.onnx> \
|
||||
--gemma=<path-to-gemma2_3b.onnx> \
|
||||
--tokenizer_path=<path-to-opencv-tokenizer-config.json> \
|
||||
--input=<path-to-image> \
|
||||
--prompt="cap en\n"
|
||||
|
||||
The tokenizer_path should point to an OpenCV-format config.json, NOT the
|
||||
HuggingFace tokenizer_config.json.
|
||||
'''
|
||||
|
||||
import numpy as np
|
||||
import argparse
|
||||
import cv2 as cv
|
||||
|
||||
EOS_ID = 1
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Use this script to run PaliGemma2 vision-language inference in OpenCV',
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--siglip', type=str, required=True, help='Path to SigLIP vision encoder ONNX model file.')
|
||||
parser.add_argument('--embedding', type=str, required=True, help='Path to embedding ONNX model file.')
|
||||
parser.add_argument('--gemma', type=str, required=True, help='Path to Gemma2 language model ONNX model file.')
|
||||
parser.add_argument('--tokenizer_path', type=str, required=True, help='Path to tokenizer config.json.')
|
||||
parser.add_argument('--input', '-i', type=str, required=True, help='Path to the input image.')
|
||||
parser.add_argument('--prompt', type=str, default='cap en\n', help='Task prompt (e.g. "cap en\\n" to caption in English).')
|
||||
parser.add_argument('--max_new_tokens', type=int, default=64, help='Maximum number of new tokens to generate.')
|
||||
parser.add_argument('--seed', type=int, default=0, help='Random seed.')
|
||||
return parser.parse_args()
|
||||
|
||||
def preprocess_image(image_path):
|
||||
'''Resize to 224x224 and normalize to [-1, 1] in CHW order (SigLIP: mean=0.5, std=0.5).'''
|
||||
img = cv.imread(image_path)
|
||||
if img is None:
|
||||
raise IOError("Could not read image: " + image_path)
|
||||
img = cv.resize(img, (224, 224))
|
||||
img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
|
||||
img = img.astype(np.float32) / 255.0
|
||||
img = (img - 0.5) / 0.5
|
||||
img = img.transpose(2, 0, 1)[np.newaxis]
|
||||
return img
|
||||
|
||||
def vlm_inference(siglip_net, embed_net, gemma_net, pixel_values, prompt, max_new_tokens, tokenizer):
|
||||
|
||||
print("Inferencing PaliGemma2 model...")
|
||||
|
||||
tokens = list(tokenizer.encode(prompt))
|
||||
input_ids = np.array([tokens], dtype=np.int64)
|
||||
|
||||
# SigLIP vision encoder: image -> image-feature tokens
|
||||
siglip_net.setInput(pixel_values, 'pixel_values')
|
||||
image_features = siglip_net.forward() # (1, 256, 2304)
|
||||
|
||||
# Text embedding: token ids -> text embeddings
|
||||
embed_net.setInput(input_ids, 'input_ids')
|
||||
text_embeds = embed_net.forward() # (1, text_len, 2304)
|
||||
|
||||
# Combine [image_features | text_embeds]
|
||||
inputs_embeds = np.concatenate([image_features, text_embeds], axis=1)
|
||||
|
||||
generated = []
|
||||
|
||||
# Prefill
|
||||
gemma_net.setInput(inputs_embeds, 'inputs_embeds')
|
||||
logits = gemma_net.forward()
|
||||
new_id = int(np.argmax(logits[0, -1, :]))
|
||||
generated.append(new_id)
|
||||
|
||||
# Decode (no KV-cache: feed full growing sequence each step)
|
||||
for _ in range(max_new_tokens - 1):
|
||||
if new_id == EOS_ID:
|
||||
break
|
||||
embed_net.setInput(np.array([[new_id]], dtype=np.int64), 'input_ids')
|
||||
new_embed = embed_net.forward()
|
||||
inputs_embeds = np.concatenate([inputs_embeds, new_embed], axis=1)
|
||||
gemma_net.setInput(inputs_embeds, 'inputs_embeds')
|
||||
logits = gemma_net.forward()
|
||||
new_id = int(np.argmax(logits[0, -1, :]))
|
||||
generated.append(new_id)
|
||||
|
||||
if generated and generated[-1] == EOS_ID:
|
||||
generated.pop()
|
||||
|
||||
return generated
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parse_args()
|
||||
np.random.seed(args.seed)
|
||||
|
||||
print("Preparing PaliGemma2 model...")
|
||||
tokenizer = cv.dnn.Tokenizer.load(args.tokenizer_path)
|
||||
|
||||
siglip_net = cv.dnn.readNetFromONNX(args.siglip, cv.dnn.ENGINE_NEW)
|
||||
embed_net = cv.dnn.readNetFromONNX(args.embedding, cv.dnn.ENGINE_NEW)
|
||||
gemma_net = cv.dnn.readNetFromONNX(args.gemma, cv.dnn.ENGINE_NEW)
|
||||
|
||||
print(f"Prompt:\n{args.prompt}")
|
||||
pixel_values = preprocess_image(args.input)
|
||||
|
||||
generated = vlm_inference(siglip_net, embed_net, gemma_net, pixel_values,
|
||||
args.prompt, args.max_new_tokens, tokenizer)
|
||||
response = tokenizer.decode(generated)
|
||||
print(f"Response:\n{response}")
|
||||
Reference in New Issue
Block a user