Compare commits
8
Commits
b92b867583
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b8b4f10ee2 | ||
|
|
0e04debb7b | ||
|
|
47cfade856 | ||
|
|
0a07cffe2a | ||
|
|
b5a0d5f1f3 | ||
|
|
8f9aa11e99 | ||
|
|
2f0c97c9bd | ||
|
|
2ff6b85d56 |
@@ -1,3 +1,8 @@
|
||||
.idea
|
||||
*/img
|
||||
*/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`
|
||||
|
||||
## 基于PyTorch的Minist数据集分类
|
||||
## **实验二** 基于PyTorch的Minist数据集分类
|
||||
|
||||
### 数据预览测试和数据装载;
|
||||
|
||||
@@ -45,3 +45,13 @@
|
||||
### 基于pytorch卷积神经网络的模型测试。
|
||||
|
||||
`Lab-2/src/model_training.py`
|
||||
|
||||
## **实验三** 迁移resnet50猫狗图像分类
|
||||
|
||||
### 通过代码自动下载模型并直接调用;对当前迁移过来的模型进行全连接层的调整;模型训练及测试结果。
|
||||
|
||||
`Lab-3/src/train.py`
|
||||
|
||||
## **实验四** 基于YOLO的目标检测
|
||||
|
||||
## **实验五** 基于YOLOv8的工业缺陷检测实验
|
||||
|
||||
Reference in New Issue
Block a user