haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Modify get_rect_adapt_landmark to adaptively decode any number of key points abd repair gen_wts.py (#1514)

Browse Source 
* Add the generation of multi-class pose engines

* Modify get_rect_adapt_landmark to adaptively decode any number of key points abd repair gen_wts.py

---------

Co-authored-by: lindsayshuo <lindsayshuo@foxmail.com.com>
This commit is contained in:
lindsayshuo 2024-05-07 19:13:15 +08:00 committed by GitHub
parent f2e8b2bc62
commit f201451f85
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
yolov8/gen_wts.py
View File

@ -1,4 +1,4 @@
import sys
import sys # noqa: F401
import argparse
import os
import struct
@ -12,7 +12,7 @@ def parse_args():
parser.add_argument(
'-o', '--output', help='Output (.wts) file path (optional)')
parser.add_argument(
'-t', '--type', type=str, default='detect', choices=['detect', 'cls', 'seg'],
'-t', '--type', type=str, default='detect', choices=['detect', 'cls', 'seg', 'pose'],
help='determines the model is detection/classification')
args = parser.parse_args()
if not os.path.isfile(args.weights):
@ -39,7 +39,7 @@ device = 'cpu'
# Load model
model = torch.load(pt_file, map_location=device)['model'].float() # load to FP32
if m_type in ['detect', 'seg']:
if m_type in ['detect', 'seg', 'pose']:
anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
delattr(model.model[-1], 'anchors')

8
yolov8/src/postprocess.cpp
View File

@ -28,7 +28,7 @@ cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
return cv::Rect(round(l), round(t), round(r - l), round(b - t));
}
cv::Rect get_rect_adapt_landmark(cv::Mat& img, float bbox[4], float lmk[51]) {
cv::Rect get_rect_adapt_landmark(cv::Mat& img, float bbox[4], float lmk[kNumberOfPoints * 3]) {
int l, r, t, b;
float r_w = kInputW / (img.cols * 1.0);
float r_h = kInputH / (img.rows * 1.0);
@ -37,7 +37,7 @@ cv::Rect get_rect_adapt_landmark(cv::Mat& img, float bbox[4], float lmk[51]) {
r = bbox[2] / r_w;
t = (bbox[1] - (kInputH - r_w * img.rows) / 2) / r_w;
b = (bbox[3] - (kInputH - r_w * img.rows) / 2) / r_w;
for (int i = 0; i < 51; i += 3) {
for (int i = 0; i < kNumberOfPoints * 3; i += 3) {
lmk[i] /= r_w;
lmk[i + 1] = (lmk[i + 1] - (kInputH - r_w * img.rows) / 2) / r_w;
// lmk[i + 2]
@ -47,7 +47,7 @@ cv::Rect get_rect_adapt_landmark(cv::Mat& img, float bbox[4], float lmk[51]) {
r = (bbox[2] - (kInputW - r_h * img.cols) / 2) / r_h;
t = bbox[1] / r_h;
b = bbox[3] / r_h;
for (int i = 0; i < 51; i += 3) {
for (int i = 0; i < kNumberOfPoints * 3; i += 3) {
lmk[i] = (lmk[i] - (kInputW - r_h * img.cols) / 2) / r_h;
lmk[i + 1] /= r_h;
// lmk[i + 2]
@ -172,7 +172,7 @@ void draw_bbox_keypoints_line(std::vector<cv::Mat>& img_batch, std::vector<std::
cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2,
cv::Scalar(0xFF, 0xFF, 0xFF), 2);
for (int k = 0; k < 51; k += 3) {
for (int k = 0; k < kNumberOfPoints * 3; k += 3) {
if (res[j].keypoints[k + 2] > 0.5) {
cv::circle(img, cv::Point((int)res[j].keypoints[k], (int)res[j].keypoints[k + 1]), 3,
cv::Scalar(0, 0x27, 0xC1), -1);