From 1696819abbea7868ae831cce15f483fb1236cc5b Mon Sep 17 00:00:00 2001 From: Gursimar Singh Date: Thu, 17 Oct 2024 13:16:53 +0530 Subject: [PATCH] Merge pull request #25667 from gursimarsingh:improved_person_reid_python 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 --- samples/dnn/models.yml | 24 ++ samples/dnn/person_reid.cpp | 549 ++++++++++++++++++++++-------------- samples/dnn/person_reid.py | 426 +++++++++++++++------------- 3 files changed, 588 insertions(+), 411 deletions(-) diff --git a/samples/dnn/models.yml b/samples/dnn/models.yml index c8a2bca3a8..23b79ce0d5 100644 --- a/samples/dnn/models.yml +++ b/samples/dnn/models.yml @@ -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" diff --git a/samples/dnn/person_reid.cpp b/samples/dnn/person_reid.cpp index 0199db1566..506e7446f5 100644 --- a/samples/dnn/person_reid.cpp +++ b/samples/dnn/person_reid.cpp @@ -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 -// Feng Zheng -// Xinyang Jiang -// Fufu Yu -// Enwei Zhang -// -// 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 + Feng Zheng + Xinyang Jiang + Fufu Yu + Enwei Zhang + +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 #include @@ -19,227 +33,340 @@ #include #include +#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(y,x)[c] = (float)((img.at(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 normalization(const std::vector& feature) -{ - std::vector 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& imglist, Net* net, const int& batch_size, const int& resize_h, const int& resize_w, std::vector>& features) -{ - for(int st = 0; st < (int)imglist.size(); st += batch_size) - { - std::vector 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 temp_feature; - for(int j = 0; j < (int)out.size().width; j++) - { - temp_feature.push_back(out.at(i,j)); - } - features.push_back(normalization(temp_feature)); - } - } - return ; -} - -static void getNames(const std::string& ImageList, std::vector& 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& feature1, const std::vector& 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>& queryFeatures, const std::vector>& galleryFeatures, const int& topk, std::vector>& result) -{ - for(int i = 0; i < (int)queryFeatures.size(); i++) - { - std::vector similarityList; - std::vector 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 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& galleryImageNames, const std::vector& 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 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>& topk, const std::vector& queryImageNames, const std::vector& 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 imgDict; +int height, width, yoloHeight, yoloWidth; +float scale, yoloScale; +bool swapRB, yoloSwapRB; +Scalar mean_v, stnd; + + +static void extractFeatures(vector &imglist, Net &net, vector &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 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 &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(queryFeatures.dot(galleryFeatures[j])); + if (currentSimilarity > maxSimilarity) + { + maxSimilarity = currentSimilarity; + bestIndex = j; + } + } + return bestIndex; +} + +static void yoloDetector(Mat &frame, Net &net, vector& 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(length) / yoloWidth; + + Mat blob; + blobFromImage(image, blob, yoloScale, Size(yoloWidth, yoloHeight), Scalar(), yoloSwapRB, false, CV_32F); + net.setInput(blob); + + vector 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 boxes; + vector scores; + vector 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(i, 0); + double centerY = outputTransposed.at(i, 1); + double w = outputTransposed.at(i, 2); + double h = outputTransposed.at(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(maxScore)); + class_ids.push_back(maxClassLoc.x); // x location gives the index + } + } + + // Apply Non-Maximum Suppression + vector 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(round(boxes[index].x * norm_scale)); + int y = static_cast(round(boxes[index].y * norm_scale)); + int w = static_cast(round(boxes[index].width * norm_scale)); + int h = static_cast(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<("@alias"); + string zooFile = findFile(parser.get("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("model"); - const std::string queryImageList = parser.get("query_list"); - const std::string galleryImageList = parser.get("gallery_list"); - const int backend = parser.get("backend"); - const int target = parser.get("target"); - const int batch_size = parser.get("batch_size"); - const int resize_h = parser.get("resize_h"); - const int resize_w = parser.get("resize_w"); - const int topk = parser.get("topk"); - const std::string output_dir= parser.get("output_dir"); + const string sha1 = parser.get("sha1"); + const string yoloSha1 = parser.get("yolo_sha1"); + const string modelPath = findModel(parser.get("model"), sha1); + const string queryImagePath = parser.get("query"); + string videoPath = parser.get("input"); + const string yoloPath = findModel(parser.get("yolo_model"), yoloSha1); + const string backend = parser.get("backend"); + const string target = parser.get("target"); + height = parser.get("height"); + width = parser.get("width"); + yoloHeight = parser.get("yolo_height"); + yoloWidth = parser.get("yolo_width"); + scale = parser.get("scale"); + yoloScale = parser.get("yolo_scale"); + swapRB = parser.get("rgb"); + yoloSwapRB = parser.get("yolo_rgb"); + mean_v = parser.get("mean"); + stnd = parser.get("std"); + int stdSize = 20; + int stdWeight = 400; + int stdImgSize = 512; + int imgWidth = -1; // Initialization + int fontSize = 50; + int fontWeight = 500; - std::vector queryImageNames; - reid::getNames(queryImageList, queryImageNames); - std::vector 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."<> queryFeatures; - reid::extractFeatures(queryImageNames, &net, batch_size, resize_h, resize_w, queryFeatures); - std::vector> galleryFeatures; - reid::extractFeatures(galleryImageNames, &net, batch_size, resize_h, resize_w, galleryFeatures); + FontFace fontFace("sans"); - std::vector> 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 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 queryFeatures; + extractFeatures(queryImages, reidNet, queryFeatures); + + vector detectedImages; + vector 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(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; } diff --git a/samples/dnn/person_reid.py b/samples/dnn/person_reid.py index 08f04faa52..48fd0d28e6 100644 --- a/samples/dnn/person_reid.py +++ b/samples/dnn/person_reid.py @@ -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 @@ -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() \ No newline at end of file