流水账记录一下yolo目标检测
1.搭建pytorch 不做解释 看以往博客或网上搜都行
2.下载yolo源码 ： https://github.com/ultralytics/ultralytics
3.样本标注工具：labelme 自己下载
4.准备数据集
4.1 新建一个放置数据集的路径
4.2 构建训练集和测试集
运行以下脚本，将数据集划分为训练集和测试集，比例是7：3，看自己需求
import os
import shutil
from tqdm import tqdm
import random
""" 使用：只需要修改 1. Dataset_folde，
2. os.chdir(os.path.join(Dataset_folder, 'images'))里的 images,
3. val_scal = 0.2
4. os.chdir('../label_json') label_json换成自己json标签文件夹名称 """
# 图片文件夹与json标签文件夹的根目录
Dataset_folder = r'D:\Software\Python\deeplearing\PytorchTest\ultralytics-main\datasets\PepperGrab'
# 把当前工作目录改为指定路径
os.chdir(os.path.join(Dataset_folder, 'images')) # images ： 图片文件夹的名称
folder = '.' # 代表os.chdir(os.path.join(Dataset_folder, 'images'))这个路径
imgs_list = os.listdir(folder)
random.seed(123) # 固定随机种子，防止运行时出现bug后再次运行导致imgs_list 里面的图片名称顺序不一致
random.shuffle(imgs_list) # 打乱
val_scal = 0.3 # 验证集比列
val_number = int(len(imgs_list) * val_scal)
val_files = imgs_list[:val_number]
train_files = imgs_list[val_number:]
print('all_files:', len(imgs_list))
print('train_files:', len(train_files))
print('val_files:', len(val_files))
os.mkdir('train')
for each in tqdm(train_files):
shutil.move(each, 'train')
os.mkdir('val')
for each in tqdm(val_files):
shutil.move(each, 'val')
os.chdir('../label_json')
os.mkdir('train')
for each in tqdm(train_files):
json_file = os.path.splitext(each)[0] + '.json'
shutil.move(json_file, 'train')
os.mkdir('val')
for each in tqdm(val_files):
json_file = os.path.splitext(each)[0] + '.json'
shutil.move(json_file, 'val')
print('划分完成')
4.2 然后就开始标样本了，用labelme标，本次测试出了目标检测外，还需要检测关键点，类似标成下面这种，我之前标过一次样本，图就不展示了，不安全.3 因为labelme标注完是json格式，要转化为yolo格式，运行以下脚本
import os
import json
import shutil
import numpy as np
from tqdm import tqdm
""""#使用：1.改 bbox_class = {'sjb_rect': 0}，我的框的类别是sjb_rect，赋值它为0，如你是dog则改成：bbox_cls = {'dog': 0}
2.改 keypoint_class = ['angle_30', 'angle_60', 'angle_90']，我的关键点类别是三个，分别是'angle_30', 'angle_60', 'angle_90'
3.改 Dataset_root 成你自己的图片与json文件的路径
4.改 os.chdir('json_label/train')与os.chdir('json_label/val') 成你的json文件夹下的train与val文件夹 """
# 数据集根慕录(即图片文件夹与标签文件夹的上一级目录)
Dataset_root = r'D:\Software\Python\deeplearing\PytorchTest\ultralytics-main\datasets\PepperGrab'
# 框的类别
bbox_class = {
'pepper': 0
}
# 关键点的类别
keypoint_class = ['head', 'end'] # 这里类别放的顺序对应关键点类别的标签 0，1，2
os.chdir(Dataset_root)
os.mkdir('labels')
os.mkdir('labels/train')
os.mkdir('labels/val')
def process_single_json(labelme_path, save_folder='../../labels/train'):
with open(labelme_path, 'r', encoding='utf-8') as f:
labelme = json.load(f)
img_width = labelme['imageWidth'] # 图像宽度
img_height = labelme['imageHeight'] # 图像高度
# 生成 YOLO 格式的 txt 文件
suffix = labelme_path.split('.')[-2]
yolo_txt_path = suffix + '.txt'
with open(yolo_txt_path, 'w', encoding='utf-8') as f:
for each_ann in labelme['shapes']: # 遍历每个标注
if each_ann['shape_type'] == 'rectangle': # 每个框，在 txt 里写一行
yolo_str = ''
## 框的信息
# 框的类别 ID
bbox_class_id = bbox_class[each_ann['label']]
yolo_str += '{} '.format(bbox_class_id)
# 左上角和右下角的 XY 像素坐标
bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))
bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))
bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))
bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))
# 框中心点的 XY 像素坐标
bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)
bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)
# 框宽度
bbox_width = bbox_bottom_right_x - bbox_top_left_x
# 框高度
bbox_height = bbox_bottom_right_y - bbox_top_left_y
# 框中心点归一化坐标
bbox_center_x_norm = bbox_center_x / img_width
bbox_center_y_norm = bbox_center_y / img_height
# 框归一化宽度
bbox_width_norm = bbox_width / img_width
# 框归一化高度
bbox_height_norm = bbox_height / img_height
yolo_str += '{:.5f} {:.5f} {:.5f} {:.5f} '.format(bbox_center_x_norm, bbox_center_y_norm,
bbox_width_norm, bbox_height_norm)
## 找到该框中所有关键点，存在字典 bbox_keypoints_dict 中
bbox_keypoints_dict = {}
for each_ann in labelme['shapes']: # 遍历所有标注
if each_ann['shape_type'] == 'point': # 筛选出关键点标注
# 关键点XY坐标、类别
x = int(each_ann['points'][0][0])
y = int(each_ann['points'][0][1])
label = each_ann['label']
if (x > bbox_top_left_x) & (x < bbox_bottom_right_x) & (y < bbox_bottom_right_y) & (
y > bbox_top_left_y): # 筛选出在该个体框中的关键点
bbox_keypoints_dict[label] = [x, y]
## 把关键点按顺序排好
for each_class in keypoint_class: # 遍历每一类关键点
if each_class in bbox_keypoints_dict:
keypoint_x_norm = bbox_keypoints_dict[each_class][0] / img_width
keypoint_y_norm = bbox_keypoints_dict[each_class][1] / img_height
yolo_str += '{:.5f} {:.5f} {} '.format(keypoint_x_norm, keypoint_y_norm,
2) # 2-可见不遮挡 1-遮挡 0-没有点
else: # 不存在的点，一律为0
yolo_str += '0 0 0 '
# 写入 txt 文件中
f.write(yolo_str + '\n')
shutil.move(yolo_txt_path, save_folder)
print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))
os.chdir('label_json/train')
save_folder = '../../labels/train'
for labelme_path in os.listdir():
try:
process_single_json(labelme_path, save_folder=save_folder)
except:
print('******有误******', labelme_path)
print('YOLO格式的txt标注文件已保存至 ', save_folder)
os.chdir('../../')
os.chdir('label_json/val')
save_folder = '../../labels/val'
for labelme_path in os.listdir():
try:
process_single_json(labelme_path, save_folder=save_folder)
except:
print('******有误******', labelme_path)
print('YOLO格式的txt标注文件已保存至 ', save_folder)
os.chdir('../../')
os.chdir('../')
4.3 构造yaml
直接拷贝…\ultralytics-main\ultralytics\cfg\models\11\yolo11-pose.yaml
修改内容如下：
4.4 开始训练模型
先下载预训练模型 yolo11n.pt yolo11n-pose.pt
然后直接训练，先不看详细训练参数，先能跑起来
from ultralytics import YOLO
import cv2
# #训练
model = YOLO("./yolo11n-pose.pt")
model.train(data = "...../ultralytics-main_0829/ultralytics-main/yolo11-pose.yaml",workers=0,epochs=640,batch=8)
跑起来了
4.5 开始预测结果
yolo = YOLO("best.pt", task = "detect")
result = yolo(source=".../ultralytics-main_0829/ultralytics-main/datasets/PepperGrab/images/val",conf=0.4,vid_stride=1,iou=0.3,save = True)
4.6 移植到C++测试 续-0908
libtorch 结合着其他几个博客一起看，我这边直接贴出代码
介绍下Tensor output输出格式
第一列是我自己打印的框编号，不算
第二~五列，框中心坐标+宽高
第六七列，我有两类，所以是两类的概率
第八~十列，为关键点XY的坐标和概率
如果换其他工作场景，需要对应更改代码，类别如果多了，那蓝色框列数增加，如果关键点多了，光色框对应增加（多一个点+3列）
C++推理代码
#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <torch/torch.h>
#include <torch/script.h>
using torch::indexing::Slice;
using torch::indexing::None;
struct ResizeParams
{
//实际图像的左上角点与原始图像0 0 的偏移
int top;
int left;
//去除灰边的图像宽高
int imgw;
int imgh;
//ratio 缩放图像和原始图像的宽高比例
float ratiow;
float ratioh;
};
struct KeyPoint
{
cv::Point2f p;
float prob;
};
struct Object
{
std::vector<KeyPoint> keypoints;
cv::Rect_<float> rect;
int label;
float prob;
};
float intersection_area(const Object& a, const Object& b)
{
cv::Rect_<float> inter = a.rect & b.rect;
return inter.area();
}
//=====Polygon_NMS=====//
void nms_sorted_bboxesPose(const std::vector<Object>& objects, std::vector<int>& picked, float nms_threshold, bool agnostic = false)
{
picked.clear();
const int n = objects.size();
std::vector<float> areas(n);
for (int i = 0; i < n; i++)
{
areas[i] = objects[i].rect.area();
}
for (int i = 0; i < n; i++)
{
const Object& a = objects[i];
int keep = 1;
for (int j = 0; j < (int)picked.size(); j++)
{
const Object& b = objects[picked[j]];
if (!agnostic && a.label != b.label)
continue;
// intersection over union
float inter_area = intersection_area(a, b);
float union_area = areas[i] + areas[picked[j]] - inter_area;
// float IoU = inter_area / union_area
if (inter_area / union_area > nms_threshold)
keep = 0;
}
if (keep)
picked.push_back(i);
}
}
//flg true/false 分别为降序/升序
template<typename T>void vector_sort(std::vector<T>vector_input, std::vector<size_t>&idx,float flg){
int num = vector_input.size();
idx.resize(num);
for (int i = 0; i < num; i++)
{
idx[i] = i;
}
if(flg){
std::sort(idx.begin(), idx.end(), [&vector_input](size_t i1, size_t i2){return vector_input[i1] > vector_input[i2]; });
}else{
std::sort(idx.begin(), idx.end(), [&vector_input](size_t i1, size_t i2){return vector_input[i1] < vector_input[i2]; });
}
}
//开始非极大抑制前的数据准备
void PrepareData(torch::Tensor output, std::vector<Object>& objects,int boxlabelnum,int keypointnum,float conf) {
objects.clear();
auto output_data = output.squeeze();
int num_params = output_data.size(0); // 11
int num_boxes = output_data.size(1); // 8400
int classnum = num_params - 5;
std::vector<Object>Tobjvec;
std::vector<float>confvec;
for (int i = 0; i < num_boxes; i++) {
float x_center = output_data[0][i].item<float>(); // 中心 x 坐标
float y_center = output_data[1][i].item<float>(); // 中心 y 坐标
float width = output_data[2][i].item<float>(); // 宽度
float height = output_data[3][i].item<float>(); // 高度
//中间几项是描框的概率值
int label = 0;
float maxconf = 0;
int endid = 0;
for (int j = 0; j < boxlabelnum; j++) {
float conf_ = output_data[j+4][i].item<float>();
if (conf_ > maxconf) {
maxconf = conf_;
label = j;
endid = j + 4;
}
else {
endid = j + 4;
}
}
//如果最大的置信度小于给定的阈值，跳过
if (maxconf < conf) {
continue;
}
Object obj1;
obj1.label = label;
obj1.prob = maxconf;
obj1.rect.x = x_center - width/2;
obj1.rect.y = y_center - height/2;
obj1.rect.width = width;
obj1.rect.height = height;
//赋值关键点
for (int j = endid+1; j < num_params; j+=3)
{
KeyPoint pp;
pp.p.x = output_data[j][i].item<float>();
pp.p.y = output_data[j+1][i].item<float>();
pp.prob = output_data[j+2][i].item<float>();
obj1.keypoints.push_back(pp);
}
Tobjvec.push_back(obj1);
confvec.push_back(maxconf);
}
//按照概率大小进行排序，降序
std::vector<size_t>sortid;
vector_sort(confvec, sortid, true);
for (int i = 0; i < sortid.size(); i++) {
objects.push_back(Tobjvec[sortid[i]]);
}
}
float generate_scale(cv::Mat& image, const std::vector<int>& target_size) {
int origin_w = image.cols;
int origin_h = image.rows;
int target_h = target_size[0];
int target_w = target_size[1];
float ratio_h = static_cast<float>(target_h) / static_cast<float>(origin_h);
float ratio_w = static_cast<float>(target_w) / static_cast<float>(origin_w);
float resize_scale = std::min(ratio_h, ratio_w);
return resize_scale;
}
float letterbox(cv::Mat &input_image, cv::Mat &output_image, const std::vector<int> &target_size, ResizeParams ¶ms) {
if (input_image.cols == target_size[1] && input_image.rows == target_size[0]) {
if (input_image.data == output_image.data)
{
return 1.;
}
else
{
output_image = input_image.clone();
return 1.;
}
}
float resize_scale = generate_scale(input_image, target_size);
int new_shape_w = std::round(input_image.cols * resize_scale);
int new_shape_h = std::round(input_image.rows * resize_scale);
float padw = (target_size[1] - new_shape_w) / 2.;
float padh = (target_size[0] - new_shape_h) / 2.;
int top = std::round(padh - 0.1);
int bottom = std::round(padh + 0.1);
int left = std::round(padw - 0.1);
int right = std::round(padw + 0.1);
params.top = top;
params.left = left;
params.imgw = new_shape_w;
params.imgh = new_shape_h;
params.ratiow = resize_scale;
params.ratioh = resize_scale;
cv::resize(input_image, output_image,
cv::Size(new_shape_w, new_shape_h),
0, 0, cv::INTER_AREA);
cv::copyMakeBorder(output_image, output_image, top, bottom, left, right,
cv::BORDER_CONSTANT, cv::Scalar(114., 114., 114));
return resize_scale;
}
void drawImage(cv::Mat &image, Object obj1, ResizeParams params)
{
std::vector<cv::Scalar>colorlabel;
colorlabel.push_back(cv::Scalar(255, 0, 0));
colorlabel.push_back(cv::Scalar(0, 255, 0));
colorlabel.push_back(cv::Scalar(0, 0, 255));
//四角点坐标
cv::Point2f ori_corners[4];
ori_corners[0].x = obj1.rect.x;
ori_corners[0].y = obj1.rect.y;
ori_corners[1].x = obj1.rect.x + obj1.rect.width;
ori_corners[1].y = obj1.rect.y;
ori_corners[2].x = obj1.rect.x + obj1.rect.width;
ori_corners[2].y = obj1.rect.y + obj1.rect.height;
ori_corners[3].x = obj1.rect.x;
ori_corners[3].y = obj1.rect.y + obj1.rect.height;
//换算回原始图像
for (size_t i = 0; i < 4; i++) {
ori_corners[i].x = (ori_corners[i].x - params.left) / params.ratiow;
ori_corners[i].y = (ori_corners[i].y - params.top) / params.ratioh;
}
for (int i = 0; i < 4; i++) {
cv::line(image, ori_corners[i], ori_corners[(i + 1) % 4], colorlabel[obj1.label], 1);
}
//开始换算关键点
std::vector<cv::Point2f> pvec;
for (int i = 0; i < obj1.keypoints.size(); i++) {
pvec.push_back(obj1.keypoints[i].p);
}
for (size_t i = 0; i < pvec.size(); i++) {
pvec[i].x = (pvec[i].x - params.left) / params.ratiow;
pvec[i].y = (pvec[i].y - params.top) / params.ratioh;
}
//如果点数过多，需要扩充colorlabel，否则段错误
for (int i = 0; i < pvec.size(); i++) {
cv::circle(image, pvec[i], 2, colorlabel[i], -1); // 红色圆点
}
}
std::vector<float>RotatedRectToXY4(Object obj1, ResizeParams params) {
//四角点坐标
cv::Point2f ori_corners[4];
ori_corners[0].x = obj1.rect.x;
ori_corners[0].y = obj1.rect.y;
ori_corners[1].x = obj1.rect.x + obj1.rect.width;
ori_corners[1].y = obj1.rect.y;
ori_corners[2].x = obj1.rect.x + obj1.rect.width;
ori_corners[2].y = obj1.rect.y + obj1.rect.height;
ori_corners[3].x = obj1.rect.x;
ori_corners[3].y = obj1.rect.y + obj1.rect.height;
//换算回原始图像
for (size_t i = 0; i < 4; i++) {
ori_corners[i].x = (ori_corners[i].x - params.left) / params.ratiow;
ori_corners[i].y = (ori_corners[i].y - params.top) / params.ratioh;
}
//开始换算关键点
std::vector<cv::Point2f> pvec;
for (int i = 0; i < obj1.keypoints.size(); i++) {
pvec.push_back(obj1.keypoints[i].p);
}
for (size_t i = 0; i < pvec.size(); i++) {
pvec[i].x = (pvec[i].x - params.left) / params.ratiow;
pvec[i].y = (pvec[i].y - params.top) / params.ratioh;
}
//输出的结果
std::vector<float>output;
for (size_t i = 0; i < 4; i++) {
output.push_back(ori_corners[i].x);
output.push_back(ori_corners[i].y);
}
output.push_back(obj1.label);
output.push_back(obj1.prob);
for (size_t i = 0; i < pvec.size(); i++)
{
output.push_back(pvec[i].x);
output.push_back(pvec[i].y);
}
return output;
}
#include<windows.h>
int main() {
printf("torch::cuda::is_available:%d\n", torch::cuda::is_available());
// Device
//torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU);
torch::Device device(torch::kCUDA);
//std::vector<std::string> classes{ "pepper" };
try {
//std::string model_path = ".../ultralytics-main/best_peper_pose.torchscript";
std::string model_path = ".../ultralytics-main_0829/ultralytics-main/best-pose.torchscript";
torch::jit::script::Module yolo_model;
LoadLibraryA("ATen_cuda.dll");
LoadLibraryA("c10_cuda.dll");
LoadLibraryA("torch_cuda.dll");
LoadLibraryA("torchvision.dll");
yolo_model = torch::jit::load(model_path, device);
yolo_model.eval();
yolo_model.to(device, torch::kFloat32);
ResizeParams params;
// Load image and preprocess
cv::Mat image = cv::imread(".../ultralytics-main_0829/ultralytics-main/20250829-102346647.jpg");
cv::Mat input_image;
letterbox(image, input_image, { 640, 640 }, params);
cv::cvtColor(input_image, input_image, cv::COLOR_BGR2RGB);
//cv::imwrite("C:/Users/lixia/Desktop/20250829-102331640_out_resize.jpg", input_image);
torch::Tensor image_tensor = torch::from_blob(input_image.data, { input_image.rows, input_image.cols, 3 }, torch::kByte).to(device);
image_tensor = image_tensor.toType(torch::kFloat32).div(255);
image_tensor = image_tensor.permute({ 2, 0, 1 });
image_tensor = image_tensor.unsqueeze(0);
std::vector<torch::jit::IValue> inputs{ image_tensor };
auto start = std::chrono::high_resolution_clock::now();
// Inference
torch::Tensor output = yolo_model.forward(inputs).toTensor().cpu();
std::cout <<"output size = " << output.size(0) << std::endl;
std::cout << "output size = " << output.size(1) << std::endl;
std::cout << "output size = " << output.size(2) << std::endl;
auto output_data = output.squeeze();
for (int i = 0; i < output.size(2); i++) {
if (output_data[4][i].item<float>() > 0.6) {
printf("%d ", i);
for (int j = 0; j < output.size(1); j++) {
printf("%.3f ", output_data[j][i].item<float>());
}
printf("\n");
}
}
std::vector<Object> objects;
PrepareData(output, objects,2,1, 0.6);
for (int i = 0; i < objects.size(); i++) {
printf("%d %d %.3f\n",i, objects[i].label, objects[i].prob);
}
//开始非极大抑制
std::vector<int> picked;
nms_sorted_bboxesPose(objects, picked, 0.4);
for (int i = 0; i < picked.size(); i++) {
printf("%d %d %d %.3f\n", i, picked[i], objects[picked[i]].label, objects[picked[i]].prob);
}
for (int i = 0; i < picked.size(); i++) {
int id = picked[i];
//绘制框
drawImage(image, objects[id], params);
//打印四角点
std::vector<float>corner = RotatedRectToXY4(objects[id], params);
for (int j = 0; j < corner.size(); j++) {
printf("%.3f ", corner[j]);
}
printf("\n");
}
cv::imwrite(".../20250829-102331640_out.jpg", image);
}
catch (const c10::Error& e) {
std::cout << e.msg() << std::endl;
}
return 0;
}
输出结果类似这种，一个正框+一个关键点
扫一扫
5321
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
