potok/opencv
1
0
Fork 0
mirror of https://github.com/opencv/opencv.git synced 2026-07-30 07:43:03 +04:00
Code Packages Releases Wiki Activity

Merge pull request #25667 from gursimarsingh:improved_person_reid_python

Browse Source 
Improved person reid cpp and python sample #25667

#25006 

This sample has been rewritten to track a selected target in a video or camera stream. Person detection has been integrated using yolov8 and the user can provide a target image via command line or interactively select the target at start of the execution


### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [ ] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Gursimar Singh
2024-10-17 13:16:53 +05:30
committed by GitHub
parent 3cd57ea09e
commit 1696819abb
3 changed files with 588 additions and 411 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/dnn/models.yml
+24
View File
@@ -369,3 +369,27 @@ dexined:
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_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"
samples/dnn/person_reid.cpp
+338 -211
View File
@@ -1,17 +1,31 @@
//
// You can download a baseline ReID model and sample input from:
// https://github.com/ReID-Team/ReID_extra_testdata
//
// 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.
//
/*
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>
@@ -19,227 +33,340 @@
#include <opencv2/highgui.hpp>
#include <opencv2/dnn.hpp>
#include "common.hpp"
using namespace cv;
using namespace cv::dnn;
using namespace std;
std::string param_keys =
"{help h | | show help message}"
"{model m | | network model}"
"{query_list q | | list of query images}"
"{gallery_list g | | list of gallery images}"
"{batch_size | 32 | batch size of each inference}"
"{resize_h | 256 | resize input to specific height.}"
"{resize_w | 128 | resize input to specific width.}"
"{topk k | 5 | number of gallery images showed in visualization}"
"{output_dir | | path for visualization(it should be existed)}";
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;
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";
namespace cv{
namespace reid{
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}";
static Mat preprocess(const Mat& img)
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
{
const double mean[3] = {0.485, 0.456, 0.406};
const double std[3] = {0.229, 0.224, 0.225};
Mat ret = Mat(img.rows, img.cols, CV_32FC3);
for (int y = 0; y < ret.rows; y ++)
bool operator()(const Mat &a, const Mat &b) const
{
for (int x = 0; x < ret.cols; x++)
{
for (int c = 0; c < 3; c++)
{
ret.at<Vec3f>(y,x)[c] = (float)((img.at<Vec3b>(y,x)[c] / 255.0 - mean[2 - c]) / std[2 - c]);
}
}
return a.data < b.data; // This is a simple pointer comparison, not content!
}
return ret;
}
static std::vector<float> normalization(const std::vector<float>& feature)
{
std::vector<float> ret;
float sum = 0.0;
for(int i = 0; i < (int)feature.size(); i++)
{
sum += feature[i] * feature[i];
}
sum = sqrt(sum);
for(int i = 0; i < (int)feature.size(); i++)
{
ret.push_back(feature[i] / sum);
}
return ret;
}
static void extractFeatures(const std::vector<std::string>& imglist, Net* net, const int& batch_size, const int& resize_h, const int& resize_w, std::vector<std::vector<float>>& features)
{
for(int st = 0; st < (int)imglist.size(); st += batch_size)
{
std::vector<Mat> batch;
for(int delta = 0; delta < batch_size && st + delta < (int)imglist.size(); delta++)
{
Mat img = imread(imglist[st + delta]);
batch.push_back(preprocess(img));
}
Mat blob = dnn::blobFromImages(batch, 1.0, Size(resize_w, resize_h), Scalar(0.0,0.0,0.0), true, false, CV_32F);
net->setInput(blob);
Mat out = net->forward();
for(int i = 0; i < (int)out.size().height; i++)
{
std::vector<float> temp_feature;
for(int j = 0; j < (int)out.size().width; j++)
{
temp_feature.push_back(out.at<float>(i,j));
}
features.push_back(normalization(temp_feature));
}
}
return ;
}
static void getNames(const std::string& ImageList, std::vector<std::string>& result)
{
std::ifstream img_in(ImageList);
std::string img_name;
while(img_in >> img_name)
{
result.push_back(img_name);
}
return ;
}
static float similarity(const std::vector<float>& feature1, const std::vector<float>& feature2)
{
float result = 0.0;
for(int i = 0; i < (int)feature1.size(); i++)
{
result += feature1[i] * feature2[i];
}
return result;
}
static void getTopK(const std::vector<std::vector<float>>& queryFeatures, const std::vector<std::vector<float>>& galleryFeatures, const int& topk, std::vector<std::vector<int>>& result)
{
for(int i = 0; i < (int)queryFeatures.size(); i++)
{
std::vector<float> similarityList;
std::vector<int> index;
for(int j = 0; j < (int)galleryFeatures.size(); j++)
{
similarityList.push_back(similarity(queryFeatures[i], galleryFeatures[j]));
index.push_back(j);
}
sort(index.begin(), index.end(), [&](int x,int y){return similarityList[x] > similarityList[y];});
std::vector<int> topk_result;
for(int j = 0; j < min(topk, (int)index.size()); j++)
{
topk_result.push_back(index[j]);
}
result.push_back(topk_result);
}
return ;
}
static void addBorder(const Mat& img, const Scalar& color, Mat& result)
{
const int bordersize = 5;
copyMakeBorder(img, result, bordersize, bordersize, bordersize, bordersize, cv::BORDER_CONSTANT, color);
return ;
}
static void drawRankList(const std::string& queryName, const std::vector<std::string>& galleryImageNames, const std::vector<int>& topk_index, const int& resize_h, const int& resize_w, Mat& result)
{
const Size outputSize = Size(resize_w, resize_h);
Mat q_img = imread(queryName), temp_img;
resize(q_img, temp_img, outputSize);
addBorder(temp_img, Scalar(0,0,0), q_img);
putText(q_img, "Query", Point(10, 30), FONT_HERSHEY_COMPLEX, 1.0, Scalar(0,255,0), 2);
std::vector<Mat> Images;
Images.push_back(q_img);
for(int i = 0; i < (int)topk_index.size(); i++)
{
Mat g_img = imread(galleryImageNames[topk_index[i]]);
resize(g_img, temp_img, outputSize);
addBorder(temp_img, Scalar(255,255,255), g_img);
putText(g_img, "G" + std::to_string(i), Point(10, 30), FONT_HERSHEY_COMPLEX, 1.0, Scalar(0,255,0), 2);
Images.push_back(g_img);
}
hconcat(Images, result);
return ;
}
static void visualization(const std::vector<std::vector<int>>& topk, const std::vector<std::string>& queryImageNames, const std::vector<std::string>& galleryImageNames, const std::string& output_dir, const int& resize_h, const int& resize_w)
{
for(int i = 0; i < (int)queryImageNames.size(); i++)
{
Mat img;
drawRankList(queryImageNames[i], galleryImageNames, topk[i], resize_h, resize_w, img);
std::string output_path = output_dir + "/" + queryImageNames[i].substr(queryImageNames[i].rfind("/")+1);
imwrite(output_path, img);
}
return ;
}
};
};
int main(int argc, char** argv)
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);
features.resize(out.rows);
for (int i = 0; i < out.rows; i++)
{
normalize(out.row(i), features[i], 1.0, 0.0, NORM_L2);
}
}
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;
}
}
int main(int argc, char **argv)
{
// Parse command line arguments.
CommandLineParser parser(argc, argv, keys);
if (argc == 1 || parser.has("help"))
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 std::string modelPath = parser.get<String>("model");
const std::string queryImageList = parser.get<String>("query_list");
const std::string galleryImageList = parser.get<String>("gallery_list");
const int backend = parser.get<int>("backend");
const int target = parser.get<int>("target");
const int batch_size = parser.get<int>("batch_size");
const int resize_h = parser.get<int>("resize_h");
const int resize_w = parser.get<int>("resize_w");
const int topk = parser.get<int>("topk");
const std::string output_dir= parser.get<String>("output_dir");
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;
std::vector<std::string> queryImageNames;
reid::getNames(queryImageList, queryImageNames);
std::vector<std::string> galleryImageNames;
reid::getNames(galleryImageList, galleryImageNames);
Net reidNet = readNetFromONNX(modelPath);
reidNet.setPreferableBackend(getBackendID(backend));
reidNet.setPreferableTarget(getTargetID(target));
dnn::Net net = dnn::readNet(modelPath);
net.setPreferableBackend(backend);
net.setPreferableTarget(target);
if(yoloPath.empty()){
cout<<"[ERROR] Please pass path to yolov8.onnx model file using --yolo_model."<<endl;
return -1;
}
Net net = readNetFromONNX(yoloPath);
std::vector<std::vector<float>> queryFeatures;
reid::extractFeatures(queryImageNames, &net, batch_size, resize_h, resize_w, queryFeatures);
std::vector<std::vector<float>> galleryFeatures;
reid::extractFeatures(galleryImageNames, &net, batch_size, resize_h, resize_w, galleryFeatures);
FontFace fontFace("sans");
std::vector<std::vector<int>> topkResult;
reid::getTopK(queryFeatures, galleryFeatures, topk, topkResult);
reid::visualization(topkResult, queryImageNames, galleryImageNames, output_dir, resize_h, resize_w);
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);
}
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;
}
samples/dnn/person_reid.py
+226 -200
View File
@@ -1,7 +1,6 @@
#!/usr/bin/env python
'''
You can download a baseline ReID model and sample input from:
https://github.com/ReID-Team/ReID_extra_testdata
This sample detects the query person in the given video file.
Authors of samples and Youtu ReID baseline:
Xing Sun <winfredsun@tencent.com>
@@ -12,229 +11,256 @@ Authors of samples and Youtu ReID baseline:
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 *
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)
def help():
print(
'''
Use this script for Person Re-identification using OpenCV.
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)
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.
MEAN = (0.485, 0.456, 0.406)
STD = (0.229, 0.224, 0.225)
To run:
Example: python person_reid.py reid
def preprocess(images, height, width):
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):
"""
Create 4-dimensional blob from image
:param image: input image
:param height: the height of the resized input image
:param width: the width of the resized input image
"""
img_list = []
for image in images:
image = cv.resize(image, (width, height))
img_list.append(image[:, :, ::-1])
images = np.array(img_list)
images = (images / 255.0 - MEAN) / STD
input = cv.dnn.blobFromImages(images.astype(np.float32), ddepth = cv.CV_32F)
return input
def extract_feature(img_dir, model_path, batch_size = 32, resize_h = 384, resize_w = 128, backend=cv.dnn.DNN_BACKEND_OPENCV, target=cv.dnn.DNN_TARGET_CPU):
"""
Extract features from images in a target directory
:param img_dir: the input image directory
:param model_path: path to ReID model
:param batch_size: the batch size for each network inference iteration
:param resize_h: the height of the input image
:param resize_w: the width of the input image
:param backend: name of computation backend
:param target: name of computation target
Extract features from images
:param images: the input images
:param net: the model network
"""
feat_list = []
path_list = os.listdir(img_dir)
path_list = [os.path.join(img_dir, img_name) for img_name in path_list]
count = 0
# 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 i in range(0, len(path_list), batch_size):
print('Feature Extraction for images in', img_dir, 'Batch:', count, '/', len(path_list))
batch = path_list[i : min(i + batch_size, len(path_list))]
imgs = read_data(batch)
inputs = preprocess(imgs, resize_h, resize_w)
feat = run_net(inputs, model_path, backend, target)
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)
count += batch_size
feats = np.concatenate(feat_list, axis = 0)
return feats, path_list
return feats
def run_net(inputs, model_path, backend=cv.dnn.DNN_BACKEND_OPENCV, target=cv.dnn.DNN_TARGET_CPU):
def find_matching(query_feat, gallery_feat):
"""
Forword propagation for a batch of images.
:param inputs: input batch of images
:param model_path: path to ReID model
:param backend: name of computation backend
:param target: name of computation target
"""
net = cv.dnn.readNet(model_path)
net.setPreferableBackend(backend)
net.setPreferableTarget(target)
net.setInput(inputs)
out = net.forward()
out = np.reshape(out, (out.shape[0], out.shape[1]))
return out
def read_data(path_list):
"""
Read all images from a directory into a list
:param path_list: the list of image path
"""
img_list = []
for img_path in path_list:
img = cv.imread(img_path)
if img is None:
continue
img_list.append(img)
return img_list
def normalize(nparray, order=2, axis=0):
"""
Normalize a N-D numpy array along the specified axis.
:param nparry: the array of vectors to be normalized
:param order: order of the norm
:param axis: the axis of x along which to compute the vector norms
"""
norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
return nparray / (norm + np.finfo(np.float32).eps)
def similarity(array1, array2):
"""
Compute the euclidean or cosine distance of all pairs.
:param array1: numpy array with shape [m1, n]
:param array2: numpy array with shape [m2, n]
Returns:
numpy array with shape [m1, m2]
"""
array1 = normalize(array1, axis=1)
array2 = normalize(array2, axis=1)
dist = np.matmul(array1, array2.T)
return dist
def topk(query_feat, gallery_feat, topk = 5):
"""
Return the index of top K gallery images most similar to the query images
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
:param topk: number of gallery images to return
"""
sim = similarity(query_feat, gallery_feat)
index = np.argsort(-sim, axis = 1)
return [i[0:int(topk)] for i in index]
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)
def drawRankList(query_name, gallery_list, output_size = (128, 384)):
"""
Draw the rank list
:param query_name: path of the query image
:param gallery_name: path of the gallery image
"param output_size: the output size of each image in the rank list
"""
def addBorder(im, color):
bordersize = 5
border = cv.copyMakeBorder(
im,
top = bordersize,
bottom = bordersize,
left = bordersize,
right = bordersize,
borderType = cv.BORDER_CONSTANT,
value = color
)
return border
query_img = cv.imread(query_name)
query_img = cv.resize(query_img, output_size)
query_img = addBorder(query_img, [0, 0, 0])
cv.putText(query_img, 'Query', (10, 30), cv.FONT_HERSHEY_COMPLEX, 1., (0,255,0), 2)
sim = query_feat.dot(gallery_feat.T)
index = np.argmax(sim, axis=1)[0]
return index
gallery_img_list = []
for i, gallery_name in enumerate(gallery_list):
gallery_img = cv.imread(gallery_name)
gallery_img = cv.resize(gallery_img, output_size)
gallery_img = addBorder(gallery_img, [255, 255, 255])
cv.putText(gallery_img, 'G%02d'%i, (10, 30), cv.FONT_HERSHEY_COMPLEX, 1., (0,255,0), 2)
gallery_img_list.append(gallery_img)
ret = np.concatenate([query_img] + gallery_img_list, axis = 1)
return ret
def main():
if hasattr(args, 'help'):
help()
exit(1)
args.model = findModel(args.model, args.sha1)
def visualization(topk_idx, query_names, gallery_names, output_dir = 'vis'):
"""
Visualize the retrieval results with the person ReID model
:param topk_idx: the index of ranked gallery images for each query image
:param query_names: the list of paths of query images
:param gallery_names: the list of paths of gallery images
:param output_dir: the path to save the visualize results
"""
if not os.path.exists(output_dir):
os.mkdir(output_dir)
for i, idx in enumerate(topk_idx):
query_name = query_names[i]
topk_names = [gallery_names[j] for j in idx]
vis_img = drawRankList(query_name, topk_names)
output_path = os.path.join(output_dir, '%03d_%s'%(i, os.path.basename(query_name)))
cv.imwrite(output_path, vis_img)
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)
yolo_net = cv.dnn.readNet(args.yolo_model)
reid_net = cv.dnn.readNet(args.model)
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__':
parser = argparse.ArgumentParser(description='Use this script to run human parsing using JPPNet',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--query_dir', '-q', required=True, help='Path to query image.')
parser.add_argument('--gallery_dir', '-g', required=True, help='Path to gallery directory.')
parser.add_argument('--resize_h', default = 256, help='The height of the input for model inference.')
parser.add_argument('--resize_w', default = 128, help='The width of the input for model inference')
parser.add_argument('--model', '-m', default='reid.onnx', help='Path to pb model.')
parser.add_argument('--visualization_dir', default='vis', help='Path for the visualization results')
parser.add_argument('--topk', default=10, help='Number of images visualized in the rank list')
parser.add_argument('--batchsize', default=32, help='The batch size of each inference')
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 backend"% 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'
% targets)
args, _ = parser.parse_known_args()
if not os.path.isfile(args.model):
raise OSError("Model not exist")
query_feat, query_names = extract_feature(args.query_dir, args.model, args.batchsize, args.resize_h, args.resize_w, args.backend, args.target)
gallery_feat, gallery_names = extract_feature(args.gallery_dir, args.model, args.batchsize, args.resize_h, args.resize_w, args.backend, args.target)
topk_idx = topk(query_feat, gallery_feat, args.topk)
visualization(topk_idx, query_names, gallery_names, output_dir = args.visualization_dir)
args = get_args_parser()
main()