Compare commits
8
Commits
b92b867583
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b8b4f10ee2 | ||
|
|
0e04debb7b | ||
|
|
47cfade856 | ||
|
|
0a07cffe2a | ||
|
|
b5a0d5f1f3 | ||
|
|
8f9aa11e99 | ||
|
|
2f0c97c9bd | ||
|
|
2ff6b85d56 |
+6
-1
@@ -1,3 +1,8 @@
|
|||||||
.idea
|
.idea
|
||||||
*/img
|
*/img
|
||||||
*/data
|
*/data
|
||||||
|
model
|
||||||
|
runs
|
||||||
|
datasets
|
||||||
|
yolov8n.pt
|
||||||
|
dataset
|
||||||
@@ -0,0 +1,158 @@
|
|||||||
|
import torch
|
||||||
|
import torchvision
|
||||||
|
from torchvision import datasets, models, transforms
|
||||||
|
import os
|
||||||
|
from torch.autograd import Variable
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import time
|
||||||
|
|
||||||
|
model_path = '../model/model_name.pth'
|
||||||
|
model_params_path = '../model/params_name.pth'
|
||||||
|
|
||||||
|
transform = transforms.Compose(
|
||||||
|
[transforms.CenterCrop(224),
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
data_dir = "../data" # 更改为您自己的数据集路径
|
||||||
|
|
||||||
|
data_transform = {
|
||||||
|
x:transforms.Compose(
|
||||||
|
[
|
||||||
|
transforms.Resize(256), # Resize图像到256x256
|
||||||
|
transforms.CenterCrop(224), # 中心裁剪到224x224
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize(
|
||||||
|
mean=[0.485, 0.456, 0.406], # ImageNet数据集的 normalize 设置
|
||||||
|
std=[0.229, 0.224, 0.225]
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
for x in ["train","valid"]
|
||||||
|
}
|
||||||
|
|
||||||
|
image_datasets = {
|
||||||
|
x:datasets.ImageFolder(
|
||||||
|
root=os.path.join(data_dir,x), # 将输入参数中的两个名字拼接成一个完整的文件路径
|
||||||
|
transform=data_transform[x]
|
||||||
|
)
|
||||||
|
for x in ["train","valid"]
|
||||||
|
}
|
||||||
|
|
||||||
|
dataloader = {
|
||||||
|
x:torch.utils.data.DataLoader(
|
||||||
|
dataset=image_datasets[x],
|
||||||
|
batch_size=16,
|
||||||
|
shuffle=True
|
||||||
|
)
|
||||||
|
for x in ["train","valid"]
|
||||||
|
}
|
||||||
|
|
||||||
|
X_example, y_example = next(iter(dataloader["train"]))
|
||||||
|
example_classes = image_datasets["train"].classes # ['cat', 'dog']
|
||||||
|
index_classes = image_datasets["train"].class_to_idx #{'cat': 0, 'dog': 1}
|
||||||
|
|
||||||
|
Use_gpu = torch.cuda.is_available()
|
||||||
|
model = models.resnet50(pretrained=True)
|
||||||
|
print(model)
|
||||||
|
|
||||||
|
for param in model.parameters():
|
||||||
|
param.requires_grad = False
|
||||||
|
|
||||||
|
# 重新设计分类器的结构
|
||||||
|
model.fc = torch.nn.Linear(2048,2)
|
||||||
|
print(model)
|
||||||
|
|
||||||
|
if Use_gpu:
|
||||||
|
model = model.cuda()
|
||||||
|
|
||||||
|
loss_f = torch.nn.CrossEntropyLoss()
|
||||||
|
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.00001)
|
||||||
|
|
||||||
|
has_been_trained = os.path.isfile(model_path)
|
||||||
|
if has_been_trained:
|
||||||
|
epoch_n = 0
|
||||||
|
model.load_state_dict(torch.load(model_params_path))
|
||||||
|
else:
|
||||||
|
epoch_n = 1
|
||||||
|
|
||||||
|
time_open = time.time()
|
||||||
|
for epoch in range(epoch_n, epoch_n + 5): # 训练5个epoch
|
||||||
|
print("Epoch {}/{}".format(epoch, epoch_n + 4))
|
||||||
|
print("-" * 10)
|
||||||
|
|
||||||
|
for phase in ["train", "valid"]:
|
||||||
|
if phase == "train":
|
||||||
|
print("Training...")
|
||||||
|
model.train(True) # 启用 BatchNormalization 和 Dropout
|
||||||
|
else:
|
||||||
|
print("Validing...")
|
||||||
|
model.train(False) # 不启用 BatchNormalization 和 Dropout
|
||||||
|
|
||||||
|
running_loss = 0.0
|
||||||
|
running_corrects = 0
|
||||||
|
|
||||||
|
for batch, data in enumerate(dataloader[phase], 1):
|
||||||
|
X, y = data
|
||||||
|
if Use_gpu:
|
||||||
|
X, y = Variable(X.cuda()), Variable(y.cuda())
|
||||||
|
else:
|
||||||
|
X, y = Variable(X), Variable(y)
|
||||||
|
|
||||||
|
y_pred = model(X)
|
||||||
|
_, pred = torch.max(y_pred.data, 1)
|
||||||
|
|
||||||
|
optimizer.zero_grad()
|
||||||
|
loss = loss_f(y_pred, y)
|
||||||
|
|
||||||
|
if phase == "train":
|
||||||
|
loss.backward()
|
||||||
|
optimizer.step()
|
||||||
|
|
||||||
|
running_loss += loss.item()
|
||||||
|
running_corrects += torch.sum(pred == y.data)
|
||||||
|
|
||||||
|
if batch % 50 == 0 and phase == "train":
|
||||||
|
print("Batch {}, Train Loss:{:.4f},Train ACC:{:.4f}%".format(
|
||||||
|
batch, running_loss / batch, 100.0 * running_corrects / (16 * batch)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
epoch_loss = running_loss * 16 / len(image_datasets[phase])
|
||||||
|
epoch_acc = 100.0 * running_corrects / len(image_datasets[phase])
|
||||||
|
|
||||||
|
print("{} Loss:{:.4f} Acc:{:.4f}%".format(phase, epoch_loss, epoch_acc))
|
||||||
|
|
||||||
|
# 保存模型参数
|
||||||
|
best_acc = 0.0
|
||||||
|
if phase == 'valid' and epoch_acc > best_acc:
|
||||||
|
best_acc = epoch_acc
|
||||||
|
torch.save(model.state_dict(), model_params_path)
|
||||||
|
|
||||||
|
time_end = time.time() - time_open
|
||||||
|
print("程序运行时间:{}分钟...".format(int(time_end / 60)))
|
||||||
|
|
||||||
|
# 读取训练集中的一个batch
|
||||||
|
X_example, Y_example = next(iter(dataloader['train']))
|
||||||
|
|
||||||
|
if Use_gpu:
|
||||||
|
X_example, Y_example = Variable(X_example.cuda()), Variable(Y_example.cuda())
|
||||||
|
else:
|
||||||
|
X_example, Y_example = Variable(X_example), Variable(Y_example)
|
||||||
|
|
||||||
|
y_pred = model(X_example)
|
||||||
|
|
||||||
|
index_classes = image_datasets['train'].class_to_idx # 显示类别对应的索引
|
||||||
|
example_classes = image_datasets['train'].classes # 将原始图像的类别保存起来
|
||||||
|
|
||||||
|
img = torchvision.utils.make_grid(X_example)
|
||||||
|
img = img.cpu().numpy().transpose([1,2,0])
|
||||||
|
|
||||||
|
print("实际:",[example_classes[i] for i in Y_example])
|
||||||
|
_, y_pred = torch.max(y_pred.data,1)
|
||||||
|
print("预测:",[example_classes[i] for i in y_pred])
|
||||||
|
|
||||||
|
plt.imshow(img)
|
||||||
|
plt.show()
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
from ultralytics import YOLO
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
# 加载最佳模型
|
||||||
|
model = YOLO('runs/detect/yolov8n_coco128_exp1/weights/best.pt')
|
||||||
|
|
||||||
|
# 在图像上进行推理
|
||||||
|
source = 'datasets/img3.png' # 替换成你自己的测试图片路径
|
||||||
|
results = model(source, save=True, conf=0.5) # conf为置信度阈值
|
||||||
|
|
||||||
|
# 可视化结果(可选,使用OpenCV)
|
||||||
|
# 读取带标注的结果图像
|
||||||
|
result_image = results[0].plot() # 获取第一张图片的绘制结果
|
||||||
|
cv2.imshow('Detection Result', result_image)
|
||||||
|
cv2.waitKey(0)
|
||||||
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
# 也可以在视频或摄像头上进行推理
|
||||||
|
# results = model('path/to/video.mp4', save=True, show=True)
|
||||||
|
# results = model(0, show=True) # 0代表默认摄像头
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
from ultralytics import YOLO
|
||||||
|
|
||||||
|
# 1. 加载一个预训练模型(小规模模型,适合实验)
|
||||||
|
# 可选模型: YOLOv8n, YOLOv8s, YOLOv8m, YOLOv8l, YOLOv8x (nano, small, medium, large, xlarge)
|
||||||
|
model = YOLO('yolov8n.pt')
|
||||||
|
|
||||||
|
# 2. 训练模型
|
||||||
|
results = model.train(
|
||||||
|
data='coco128.yaml', # 数据集配置文件,YOLOv8会自动识别并下载
|
||||||
|
epochs=50, # 训练轮次,建议50-100轮以观察效果
|
||||||
|
imgsz=640, # 输入图像大小
|
||||||
|
batch=16, # 批次大小,根据GPU内存调整(CPU可设为-1)
|
||||||
|
name='yolov8n_coco128_exp1', # 实验名称,用于保存结果
|
||||||
|
device='cpu' # 使用CPU训练,如果有GPU且环境配置好,可改为 device=0 或 'cuda'
|
||||||
|
)
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
from ultralytics import YOLO
|
||||||
|
|
||||||
|
# 加载训练得到的最佳模型
|
||||||
|
best_model = YOLO('runs/detect/yolov8n_coco128_exp1/weights/best.pt')
|
||||||
|
# 在验证集上评估模型
|
||||||
|
metrics = best_model.val()
|
||||||
|
# 程序会自动打印出mAP、 Precision、 Recall等关键指标
|
||||||
|
|
||||||
|
# 加载训练得到的最佳模型
|
||||||
|
best_model = YOLO('runs/detect/yolov8n_coco128_exp1/weights/best.pt')
|
||||||
|
|
||||||
|
# 在验证集上评估模型
|
||||||
|
metrics = best_model.val()
|
||||||
|
# 程序会自动打印出mAP、 Precision、 Recall等关键指标。
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
from ultralytics import YOLO
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
# 加载最佳模型
|
||||||
|
model = YOLO('runs/detect/yolov8n_pcb_defect_exp14/weights/best.pt')
|
||||||
|
|
||||||
|
# 在单张图像上推理
|
||||||
|
results = model('../dataset/images/val/patches_296.jpg', save=True, conf=0.5)
|
||||||
|
|
||||||
|
# 可视化结果(使用OpenCV)
|
||||||
|
for result in results:
|
||||||
|
plotted_img = result.plot() # 绘制检测结果
|
||||||
|
cv2.imshow('Detection Result', plotted_img)
|
||||||
|
cv2.waitKey(0)
|
||||||
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
# 支持文件夹、视频、摄像头推理
|
||||||
|
# results = model('path/to/video.mp4', save=True, show=True) # 视频文件
|
||||||
|
# results = model(0, show=True) # 摄像头(0为默认摄像头)
|
||||||
|
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
path: ../dataset # 数据集根目录
|
||||||
|
train: images/train # 训练图像路径
|
||||||
|
val: images/val # 验证图像路径
|
||||||
|
|
||||||
|
# 类别数
|
||||||
|
nc: 6
|
||||||
|
# 类别名称
|
||||||
|
names: ['missing_hole', 'mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper']
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
from ultralytics import YOLO
|
||||||
|
|
||||||
|
# 加载预训练模型
|
||||||
|
model = YOLO('yolov8n.pt') # 可选模型: yolov8n, yolov8s, yolov8m, yolov8l, yolov8x
|
||||||
|
|
||||||
|
# 训练模型
|
||||||
|
results = model.train(
|
||||||
|
data='pcb_defect.yaml', # 数据集配置文件路径
|
||||||
|
epochs=20, # 训练轮次
|
||||||
|
imgsz=640, # 输入图像大小
|
||||||
|
batch=16, # 批次大小(根据GPU内存调整)
|
||||||
|
name='yolov8n_pcb_defect_exp1', # 实验名称
|
||||||
|
optimizer='AdamW', # 优化器
|
||||||
|
lr0=0.001, # 初始学习率
|
||||||
|
device='cpu' # 使用设备(如GPU:device=0)
|
||||||
|
)
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
import xml.etree.ElementTree as ET
|
||||||
|
|
||||||
|
import pickle
|
||||||
|
import os
|
||||||
|
from os import listdir, getcwd
|
||||||
|
from os.path import join
|
||||||
|
import glob
|
||||||
|
|
||||||
|
classes = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches'] # xml文件中标记的种类
|
||||||
|
|
||||||
|
|
||||||
|
def convert(size, box):
|
||||||
|
dw = 1.0 / size[0]
|
||||||
|
dh = 1.0 / size[1]
|
||||||
|
x = (box[0] + box[1]) / 2.0
|
||||||
|
y = (box[2] + box[3]) / 2.0
|
||||||
|
w = box[1] - box[0]
|
||||||
|
h = box[3] - box[2]
|
||||||
|
x = x * dw
|
||||||
|
w = w * dw
|
||||||
|
y = y * dh
|
||||||
|
h = h * dh
|
||||||
|
return (x, y, w, h)
|
||||||
|
|
||||||
|
|
||||||
|
def convert_annotation(image_name):
|
||||||
|
try:
|
||||||
|
in_file = open('../dataset/annotations/val/' + image_name[:-3] + 'xml', encoding='utf-8') # 原来的xml文件路径
|
||||||
|
except FileNotFoundError:
|
||||||
|
print(f"Warning: XML file for {image_name} not found, skipping.")
|
||||||
|
return
|
||||||
|
|
||||||
|
out_file = open('../dataset/labels/val/' + image_name[:-3] + 'txt', 'w', encoding='utf-8') # 转换后的txt文件存放路径
|
||||||
|
tree = ET.parse(in_file)
|
||||||
|
root = tree.getroot()
|
||||||
|
size = root.find('size')
|
||||||
|
w = int(size.find('width').text)
|
||||||
|
h = int(size.find('height').text)
|
||||||
|
|
||||||
|
for obj in root.iter('object'):
|
||||||
|
difficult = obj.find('difficult').text
|
||||||
|
cls = obj.find('name').text
|
||||||
|
if cls not in classes or int(difficult) == 1:
|
||||||
|
continue
|
||||||
|
cls_id = classes.index(cls)
|
||||||
|
xmlbox = obj.find('bndbox')
|
||||||
|
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
|
||||||
|
float(xmlbox.find('ymax').text))
|
||||||
|
bb = convert((w, h), b)
|
||||||
|
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
|
||||||
|
in_file.close()
|
||||||
|
out_file.close()
|
||||||
|
|
||||||
|
|
||||||
|
wd = getcwd()
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
|
||||||
|
for image_path in glob.glob("../dataset/images/val/*.jpg"): # xml对应的图片的路径
|
||||||
|
image_name = image_path.split('\\')[-1]
|
||||||
|
convert_annotation(image_name)
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
from ultralytics import YOLO
|
||||||
|
|
||||||
|
# 加载最佳模型
|
||||||
|
best_model = YOLO('runs/detect/yolov8n_pcb_defect_exp14/weights/best.pt')
|
||||||
|
|
||||||
|
# 在测试集上评估
|
||||||
|
metrics = best_model.val() # 程序会自动计算mAP等指标
|
||||||
|
print(f"mAP@0.5: {metrics.box.map50}")
|
||||||
|
print(f"mAP@0.5:0.95: {metrics.box.map}")
|
||||||
@@ -24,7 +24,7 @@
|
|||||||
|
|
||||||
`Lab-1/src/rgb_img.py`
|
`Lab-1/src/rgb_img.py`
|
||||||
|
|
||||||
## 基于PyTorch的Minist数据集分类
|
## **实验二** 基于PyTorch的Minist数据集分类
|
||||||
|
|
||||||
### 数据预览测试和数据装载;
|
### 数据预览测试和数据装载;
|
||||||
|
|
||||||
@@ -44,4 +44,14 @@
|
|||||||
|
|
||||||
### 基于pytorch卷积神经网络的模型测试。
|
### 基于pytorch卷积神经网络的模型测试。
|
||||||
|
|
||||||
`Lab-2/src/model_training.py`
|
`Lab-2/src/model_training.py`
|
||||||
|
|
||||||
|
## **实验三** 迁移resnet50猫狗图像分类
|
||||||
|
|
||||||
|
### 通过代码自动下载模型并直接调用;对当前迁移过来的模型进行全连接层的调整;模型训练及测试结果。
|
||||||
|
|
||||||
|
`Lab-3/src/train.py`
|
||||||
|
|
||||||
|
## **实验四** 基于YOLO的目标检测
|
||||||
|
|
||||||
|
## **实验五** 基于YOLOv8的工业缺陷检测实验
|
||||||
|
|||||||
Reference in New Issue
Block a user