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

添加几个脚本

Browse Source
This commit is contained in:
haotian 2025-05-21 10:51:17 +08:00
parent 7ac785ddb3
commit d78e35f837
14 changed files with 422 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
005查看yolo格式目标检测框.py
View File

@ -3,9 +3,11 @@ import cv2
import numpy as np
# 总的检测根目录
path_root_labels = 'fire-8/valid/labels'
path_root_labels = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/Safty/train/labels_raw'
# 总的检测根目录
path_root_imgs = 'fire-8/valid/images'
path_root_imgs = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/Safty/train/images'
path_root_save = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/Safty/train/show'
remove_path = 'fire-8/remove/valid.txt'
@ -34,9 +36,9 @@ for ii in os.walk(path_root_imgs):
line = f.readline()
if line:
msg = line.split(" ")
if msg[0] == '1':
with open(remove_path, 'a') as t:
t.write(j+'\n')
# if msg[0] == '1':
# with open(remove_path, 'a') as t:
# t.write(j+'\n')
# print(j)
# print(x_center,",",y_center,",",width,",",height)
x1 = int((float(msg[1]) - float(msg[3]) / 2) * w) # x_center - width/2
@ -45,9 +47,10 @@ for ii in os.walk(path_root_imgs):
y2 = int((float(msg[2]) + float(msg[4]) / 2) * h) # y_center + height/2
# print(x1, ",", y1, ",", x2, ",", y2)
cv2.rectangle(img_tmp, (x1, y1), (x2, y2), (0, 0, 255), 5)
cv2.putText(img_tmp, msg[0], (x1-5, y1-5), 2, 2, 0, 2)
else:
break
# cv2.imshow("show", img_tmp)
# c = cv2.waitKey(0)
cv2.imwrite(f'fire-8/show/{j}', img_tmp)
cv2.imwrite(f'{path_root_save}/{j}', img_tmp)
# print(f'写入文件{j}')

6
006yolov8训练.py
View File

@ -18,9 +18,9 @@ if __name__ == '__main__':
# print(torch.cuda.is_available())
model = YOLO('yolov8n.yaml')
model = YOLO('./yolov8n.yaml')
results = model.train(data='./dataset_20250304.yaml', epochs=200, batch=8, patience=20, imgsz=640, save=True,
workers=0, device=0, project='HelmetShoeHead', name='train', verbose=True)
results = model.train(data='/home/admin-root/haotian/python哈汽锻8安全帽识别/dataset_20250428_BAC.yaml', epochs=200, batch=32, patience=20, imgsz=640, save=True,
workers=0, device=0, project='HelmetHeadBAC', name='train', verbose=True)
# print(model)

42
008yolov8预测.py
View File

@ -7,26 +7,32 @@ from ultralytics import YOLO
'''
model = YOLO('5000_HeadHelmet/train2/weights/best.pt')
model = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HeadShoes/trainshoeV2/weights/best.pt')
images_path = 'test_images'
images_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images'
save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputShoeV2'
all_images = [os.path.join(images_path, t) for t in os.listdir(images_path)]
# conf 置信度在0.5 以上的框才显示。
results = model(all_images,
conf = 0.5, device=0)
# Process results list
i = 0
for result in results:
boxes = result.boxes # Boxes object for bounding box outputs
# masks = result.masks # Masks object for segmentation masks outputs
# keypoints = result.keypoints # Keypoints object for pose outputs
# 属于这个类的置信度
probs = result.probs # Probs object for classification outputs
# obb = result.obb # Oriented boxes object for OBB outputs
print(boxes,' ',probs)
# result.show() # display to screen
result.save(filename=f"result_{i}.jpg") # save to disk
i += 1
for i in range(len(all_images)):
# conf 置信度在.
results = model(all_images[i],
conf = 0.2, device=0)
# # Process results list
# = 0
# for result in results:
# boxes = result.boxes # Boxes object for bounding box outputs
# # masks = result.masks # Masks object for segmentation masks outputs
# # keypoints = result.keypoints # Keypoints object for pose outputs
# # 属于这个类的置信度
# probs = result.probs # Probs object for classification outputs
# # obb = result.obb # Oriented boxes object for OBB outputs
# print(boxes,' ',probs)
# # result.show() # display to screen
results[0].save(filename=f"{save_path}/result_{i}.jpg") # save to disk
# i += 1

4
009导出yolov8模型.py
View File

@ -1,7 +1,9 @@
from ultralytics import YOLO
'''
导出yolo模型并不能转换成tensorRT模型
'''
model = YOLO('5000_HeadHelmet/train2/weights/best.pt')
model.export(format="engine", device=0)

12
025划分数据集.py
View File

@ -4,10 +4,10 @@ import shutil
def split_dataset(images_dir, labels_dir, train_ratio=0.8):
# 创建 train 和 val 文件夹
train_images_dir = os.path.join('/home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集/train', 'images')
val_images_dir = os.path.join('/home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集/val', 'images')
train_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集/train', 'labels')
val_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集/val', 'labels')
train_images_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/train', 'images')
val_images_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/val', 'images')
train_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/train', 'labels')
val_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/val', 'labels')
os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(val_images_dir, exist_ok=True)
@ -39,8 +39,8 @@ def split_dataset(images_dir, labels_dir, train_ratio=0.8):
print(f"数据集已成功划分为训练集和验证集，训练集占比 {train_ratio * 100}%")
# 使用示例
images_dir = '/home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集/images' # 替换为你的 images 文件夹路径
labels_dir = '/home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集/labels_helmetShoeHead_filter' # 替换为你的 labels 文件夹路径
images_dir = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/images/images' # 替换为你的 images 文件夹路径
labels_dir = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/images/labels' # 替换为你的 labels 文件夹路径
train_ratio = 0.8 # 训练集占比，默认为 80%
split_dataset(images_dir, labels_dir, train_ratio)

40
026保留指定类别的标签.py Normal file
View File

@ -0,0 +1,40 @@
import os
# 类别名称列表
class_names = ['badaoxie', 'banxie', 'baomihuaxie', 'heiyaoshixie', 'huoxingxie', 'jianzhenxie', 'maozhuaxie', 'mianhuatangxie', 'mofangxie', 'moshutiexie', 'qidianxie', 'ranshaoxie', 'shuinixie', 'wangmianxie', 'xiaobaixie', 'xiaochujuxie', 'xiaomaixie', 'yingtangxie', 'yinsuxie', 'yuanyangxie', 'yucixie'] # 替换为你的类别名称
target_class = ['badaoxie', 'banxie', 'baomihuaxie', 'heiyaoshixie', 'huoxingxie', 'jianzhenxie', 'maozhuaxie', 'mianhuatangxie', 'mofangxie', 'moshutiexie', 'qidianxie', 'ranshaoxie', 'shuinixie', 'wangmianxie', 'xiaobaixie', 'xiaochujuxie', 'xiaomaixie', 'yingtangxie', 'yinsuxie', 'yuanyangxie', 'yucixie'] # 目标类别
target_class_id = {'badaoxie': 2, 'banxie':2, 'baomihuaxie':2, 'heiyaoshixie':2, 'huoxingxie':2,
'jianzhenxie':2, 'maozhuaxie':2, 'mianhuatangxie':2, 'mofangxie':2, 'moshutiexie':2,
'qidianxie':2, 'ranshaoxie':2, 'shuinixie':2, 'wangmianxie':2, 'xiaobaixie':2,
'xiaochujuxie':2, 'xiaomaixie':2, 'yingtangxie':2, 'yinsuxie':2, 'yuanyangxie':2, 'yucixie':2} # 目标类别的新 ID
def filter_and_rename_labels(label_directory, output_label_directory):
if not os.path.exists(output_label_directory):
os.makedirs(output_label_directory) # 创建输出标签目录
# 遍历标签目录中的所有文件
for label_file in os.listdir(label_directory):
if label_file.endswith('.txt'): # 确保是标签文件
with open(os.path.join(label_directory, label_file), 'r') as file:
lines = file.readlines()
# 过滤标签，只保留目标类别并修改类别 ID
filtered_lines = []
for line in lines:
try:
class_id, x_center, y_center, box_width, box_height = map(float, line.split())
if class_names[int(class_id)] in target_class:
# 修改类别 ID 为 0
filtered_lines.append(f"{target_class_id[class_names[int(class_id)]]} {x_center} {y_center} {box_width} {box_height}\n")
except:
print("标签文件错误", label_file)
# 将过滤后的标签写入新的文件
with open(os.path.join(output_label_directory, label_file), 'w') as file:
file.writelines(filtered_lines)
# 示例用法
label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/Shoes/shoe.v2i.yolov8/valid/labels_raw' # 替换为你的标签文件夹路径
output_label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/Shoes/shoe.v2i.yolov8/valid/labels' # 替换为你希望保存过滤后标签的路径
filter_and_rename_labels(label_directory, output_label_directory)

34
027保留指定类别标签_1.py Normal file
View File

@ -0,0 +1,34 @@
import os
# 类别名称列表
class_names = ['0', 'head', 'helmet', 'object', 'person'] # 替换为你的类别名称
target_class = ['head', 'helmet'] # 目标类别
def filter_labels(label_directory, output_label_directory):
if not os.path.exists(output_label_directory):
os.makedirs(output_label_directory) # 创建输出标签目录
# 遍历标签目录中的所有文件
for label_file in os.listdir(label_directory):
if label_file.endswith('.txt'): # 确保是标签文件
with open(os.path.join(label_directory, label_file), 'r') as file:
lines = file.readlines()
# 过滤标签，只保留目标类别
filtered_lines = []
for line in lines:
try:
class_id, _, _, _, _ = map(float, line.split())
if class_names[int(class_id)] in target_class:
filtered_lines.append(line)
except:
print(label_file)
# 将过滤后的标签写入新的文件
with open(os.path.join(output_label_directory, label_file), 'w') as file:
file.writelines(filtered_lines)
# 示例用法
label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集/train/labels_raw' # 替换为你的标签文件夹路径
output_label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集/train/labels_filtered' # 替换为你希望保存过滤后标签的路径
filter_labels(label_directory, output_label_directory)

49
028合并标签文件.py Normal file
View File

@ -0,0 +1,49 @@
import os
label_path_1 = '/home/admin-root/haotian/python哈汽锻8安全帽识别/头鞋数据集/head_imgs/labels_head'
label_path_2 = '/home/admin-root/haotian/python哈汽锻8安全帽识别/头鞋数据集/head_imgs/labels_shoe'
label_output_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/头鞋数据集/head_imgs/labels'
label_names = os.listdir(label_path_1)
for label_name in label_names:
# 原始标签路径, lp1 head: 1, lp2 shoe:2
lp1 = os.path.join(label_path_1, label_name)
lp2 = os.path.join(label_path_2, label_name)
# 新的标签路径
lop = os.path.join(label_output_path, label_name)
labels = list()
with open(lp1, 'r') as f:
lines1 = f.readlines()
for line in lines1:
try:
class_id, x_center, y_center, box_width, box_height = map(float, line.split())
labels.append(f"1 {x_center} {y_center} {box_width} {box_height}\n")
except:
print("lp1 解析标签文件出错")
with open(lp2, 'r') as f:
lines2 = f.readlines()
for line in lines2:
try:
class_id, x_center, y_center, box_width, box_height = map(float, line.split())
labels.append(f"2 {x_center} {y_center} {box_width} {box_height}\n")
except:
print("lp2 解析标签文件出错")
with open(lop, 'w') as f:
f.writelines(labels)

125
029将yolo预测结果写为label_studio_json格式.py Normal file
View File

@ -0,0 +1,125 @@
import os.path
import string
import json
from ultralytics import YOLO
'''
注意修改配置文件中的文件夹路径
'''
# 使用secrets模块（更安全但更慢）
import secrets
def secure_generate(num: int):
chars = string.ascii_letters + string.digits + '_'
candidate = ''.join(secrets.choice(chars) for _ in range(num))
return candidate
model = YOLO('/home/admin-root/haotian/劳保鞋识别/jingzhu_245_1112/tensorrtx-master/yolov8/build_1113/best.pt')
images_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/HardHatsV1/train/images'
save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputShoeV2'
image_names = os.listdir(images_path)
all_images = [os.path.join(images_path, t) for t in image_names]
json_list = list()
for i in range(len(all_images)):
# conf 置信度在.
results = model(all_images[i],
conf = 0.2, device=0)
# # Process results list
# = 0
# for result in results:
# boxes = result.boxes # Boxes object for bounding box outputs
# # masks = result.masks # Masks object for segmentation masks outputs
# # keypoints = result.keypoints # Keypoints object for pose outputs
# # 属于这个类的置信度
# probs = result.probs # Probs object for classification outputs
# # obb = result.obb # Oriented boxes object for OBB outputs
# print(boxes,' ',probs)
# # result.show() # display to screen
# results[0].save(filename=f"{save_path}/result_{i}.jpg") # save to disk
names = results[0].names
xywhns = results[0].boxes.xywhn
clss = results[0].boxes.cls
t_d = dict()
t_d["data"] = {"image": f"/data/local-files/?d=HardHatsV1/train/images/{image_names[i]}"}
t_d["predictions"] = list()
t_d_p = dict()
t_d_p["model_version"] = "version 1"
t_d_p["score"] = 0.8
t_d_p["result"] = list()
orig_img = results[0].orig_shape
for idx,boxe in enumerate(xywhns):
t_d_p_b = dict()
t_d_p_b["original_width"] = orig_img[1]
t_d_p_b["original_height"] = orig_img[0]
t_d_p_b["image_rotation"] = 0
#随机字符串
t_d_p_b["id"] = secure_generate(9)
t_d_p_b["from_name"] = "label"
t_d_p_b["to_name"] = "image"
t_d_p_b["type"] = "rectanglelabels"
t_d_p_b["origin"] = "manual"
x, y, w, h = boxe
class_name = names[clss[idx].item()]
t_d_p_b["value"]={
"x": x.item() * 100,
"y": y.item() * 100,
"width" : w.item() * 100,
"height" : h.item() * 100,
"rotation": 0,
"rectanglelabels": [ class_name]
}
t_d_p["result"].append(t_d_p_b)
t_d["predictions"].append(t_d_p)
json_list.append(t_d)
break
with open("test_json.json", "w", encoding="utf-8") as f:
json.dump(json_list,f,
indent=4, # 4空格缩进
sort_keys=True, # 按键名字典序排序
ensure_ascii=False # 支持非ASCII字符
)
# i += 1

128
030将yolo预测结果写为label_studio_json格式2.py Normal file
View File

@ -0,0 +1,128 @@
import json
import os
from uuid import uuid4
from ultralytics import YOLO
def yolo_to_labelstudio(results, class_names, model_version="yolov8_1.0"):
"""将 YOLOv8 检测结果转换为 Label Studio 预测格式
Args:
results: YOLOv8 Results 对象
class_names: 类别名称列表
model_version: 模型版本标识
Returns:
Label Studio 兼容的预测字典
"""
predictions = []
# 获取原始图像尺寸 (height, width)
orig_height = results[0].orig_shape[0]
orig_width = results[0].orig_shape[1]
for box in results[0].boxes:
# 提取检测框信息
x_min, y_min, x_max, y_max = box.xyxy[0].tolist()
conf = box.conf.item()
cls_idx = int(box.cls.item())
# 计算百分比坐标（保留3位小数）
x_percent = round((x_min / orig_width) * 100, 3)
y_percent = round((y_min / orig_height) * 100, 3)
width_percent = round((x_max - x_min) / orig_width * 100, 3)
height_percent = round((y_max - y_min) / orig_height * 100, 3)
# 构建预测项
prediction = {
"id": str(uuid4()), # 生成唯一ID
"type": "rectanglelabels",
"from_name": "label",
"to_name": "image",
"original_width": orig_width,
"original_height": orig_height,
"value": {
"x": x_percent,
"y": y_percent,
"width": width_percent,
"height": height_percent,
"rotation": 0,
"rectanglelabels": [class_names[cls_idx]]
},
"score": round(conf, 4) # 保留4位小数
}
predictions.append(prediction)
return {
"model_version": model_version,
"result": predictions
}
# 使用示例
if __name__ == "__main__":
# 初始化模型
# model_0 = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV1/weights/best.pt')
# 检测鞋的模型
model_1 = YOLO('/home/admin-root/haotian/劳保鞋识别/jingzhu_245_1112/tensorrtx-master/yolov8/build_1113/best.pt')
# 定义类别名称 (需与训练时一致)
CLASS_NAMES_0 = [
"helmet", "head",
]
CLASS_NAMES_1 = [
"shoe",
]
# 头,安全帽
model_0 = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV1/weights/best.pt')
images_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/HardHatsV1/train/images'
save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputShoeV2'
image_names = os.listdir(images_path)
all_images = [os.path.join(images_path, t) for t in image_names]
json_list = list()
for i in range(len(all_images)):
# conf 置信度在.
results_0 = model_0(all_images[i],
conf = 0.2, device=0)
predictions_0 = yolo_to_labelstudio(results_0, CLASS_NAMES_0)
results_1 = model_1(all_images[i],
conf=0.2, device=0)
predictions_1 = yolo_to_labelstudio(results_1, CLASS_NAMES_1)
predictions_0["result"] += predictions_1["result"]
t_d = dict()
t_d["data"] = {"image": f"/data/local-files/?d=HardHatsV1/train/images/{image_names[i]}"}
t_d["predictions"] = list()
t_d["predictions"].append(predictions_0)
json_list.append(t_d)
if i == 10:
break
with open("test_json.json", "w", encoding="utf-8") as f:
json.dump(json_list,f,
indent=4, # 4空格缩进
sort_keys=True, # 按键名字典序排序
ensure_ascii=False # 支持非ASCII字符
)

2
dataset_20250304.yaml
View File

@ -8,7 +8,7 @@
# └── coco8 ← downloads here (1 MB)
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: /home/admin-root/haotian/python哈汽安全帽识别/安全帽鞋头数据集 # dataset root dir
path: /home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽鞋头数据集 # dataset root dir
train: train # train images (relative to 'path') 4 images
val: val # val images (relative to 'path') 4 images
test: # test images (optional)

BIN
安全帽鞋头数据集.zip
View File

Binary file not shown.

BIN
安全帽鞋头数据集/train/labels.cache
View File

Binary file not shown.

BIN
安全帽鞋头数据集/val/labels.cache
View File

Binary file not shown.