Compare commits
12
Commits
1e1f9644fe
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b8b4f10ee2 | ||
|
|
0e04debb7b | ||
|
|
47cfade856 | ||
|
|
0a07cffe2a | ||
|
|
b5a0d5f1f3 | ||
|
|
8f9aa11e99 | ||
|
|
2f0c97c9bd | ||
|
|
2ff6b85d56 | ||
|
|
b92b867583 | ||
|
|
9b90328262 | ||
|
|
1b58da9a47 | ||
|
|
d9d2293999 |
@@ -1,3 +1,8 @@
|
|||||||
.idea
|
.idea
|
||||||
*/img
|
*/img
|
||||||
*/data
|
*/data
|
||||||
|
model
|
||||||
|
runs
|
||||||
|
datasets
|
||||||
|
yolov8n.pt
|
||||||
|
dataset
|
||||||
@@ -0,0 +1,67 @@
|
|||||||
|
import torch
|
||||||
|
import torchvision
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
from torchvision import transforms
|
||||||
|
from torch.utils.data import random_split, DataLoader
|
||||||
|
|
||||||
|
|
||||||
|
def load_and_preview_mnist(data_size=100, batch_size=16, seed=42, root='../data'):
|
||||||
|
"""
|
||||||
|
加载并预览 MNIST 数据集。
|
||||||
|
Parameters:
|
||||||
|
- data_size (int): 加载的数据条目数。
|
||||||
|
- batch_size (int): 每批次处理的数据数量。
|
||||||
|
- seed (int): 随机种子,用于保证数据切割的可重复性。
|
||||||
|
- root (str): 数据集保存的路径。
|
||||||
|
"""
|
||||||
|
# 数据预处理
|
||||||
|
transform = transforms.Compose([
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize((0.5,), (0.5,))
|
||||||
|
])
|
||||||
|
|
||||||
|
# 载入数据集
|
||||||
|
train_dataset = torchvision.datasets.MNIST(
|
||||||
|
root=root,
|
||||||
|
train=True,
|
||||||
|
download=True,
|
||||||
|
transform=transform
|
||||||
|
)
|
||||||
|
|
||||||
|
# 截取指定数量的数据
|
||||||
|
data_train, _ = random_split(
|
||||||
|
dataset=train_dataset,
|
||||||
|
lengths=[data_size, len(train_dataset) - data_size],
|
||||||
|
generator=torch.Generator().manual_seed(seed)
|
||||||
|
)
|
||||||
|
|
||||||
|
# 加载数据集
|
||||||
|
train_loader = DataLoader(
|
||||||
|
dataset=data_train,
|
||||||
|
batch_size=batch_size,
|
||||||
|
shuffle=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# 数据预览
|
||||||
|
images, labels = next(iter(train_loader))
|
||||||
|
|
||||||
|
# 将图像拼接为网格
|
||||||
|
img_grid = torchvision.utils.make_grid(images, nrow=batch_size // 4)
|
||||||
|
img = img_grid.numpy().transpose(1, 2, 0)
|
||||||
|
std = [0.5]
|
||||||
|
mean = [0.5]
|
||||||
|
img = img * std + mean
|
||||||
|
|
||||||
|
# 显示图像
|
||||||
|
print([labels[i] for i in range(batch_size)])
|
||||||
|
plt.imshow(img, cmap='gray')
|
||||||
|
plt.axis('off')
|
||||||
|
plt.show()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
# 使用默认参数调用函数
|
||||||
|
load_and_preview_mnist()
|
||||||
|
|
||||||
|
# 或者使用自定义参数
|
||||||
|
# load_and_preview_mnist(data_size=200, batch_size=8, seed=123, root='./data')
|
||||||
@@ -0,0 +1,89 @@
|
|||||||
|
import torch
|
||||||
|
import torchvision
|
||||||
|
import torchvision.transforms as transforms
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
from torch.utils.data import DataLoader, random_split
|
||||||
|
|
||||||
|
|
||||||
|
def get_transform():
|
||||||
|
"""
|
||||||
|
返回 torchvision.transforms.Compose 对象,包含数据预处理变换操作。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- transform (torchvision.transforms.Compose): 包含数据预处理变换的 Compose 对象。
|
||||||
|
"""
|
||||||
|
return transforms.Compose([
|
||||||
|
transforms.ToTensor(), # 将输入数据类型转换为 Tensor
|
||||||
|
transforms.Normalize((0.5,), (0.5,)) # 对数据进行标准化处理,均值和标准差均为 0.5
|
||||||
|
])
|
||||||
|
|
||||||
|
|
||||||
|
def load_mnist(data_size=100, batch_size=16, seed=42, root='../data', train=True):
|
||||||
|
"""
|
||||||
|
加载 MNIST 数据集。
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
- data_size (int): 加载的数据条目数。
|
||||||
|
- batch_size (int): 每批次处理的数据数量。
|
||||||
|
- seed (int): 随机种子,用于保证数据切割的可重复性。
|
||||||
|
- root (str): 数据集保存的路径。
|
||||||
|
- train (bool): 是否加载训练集。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- DataLoader: 数据加载器。
|
||||||
|
"""
|
||||||
|
transform = get_transform()
|
||||||
|
|
||||||
|
dataset = torchvision.datasets.MNIST(
|
||||||
|
root=root,
|
||||||
|
train=train,
|
||||||
|
download=True,
|
||||||
|
transform=transform
|
||||||
|
)
|
||||||
|
|
||||||
|
data_subset, _ = random_split(
|
||||||
|
dataset=dataset,
|
||||||
|
lengths=[data_size, len(dataset) - data_size],
|
||||||
|
generator=torch.Generator().manual_seed(seed)
|
||||||
|
)
|
||||||
|
|
||||||
|
data_loader = DataLoader(
|
||||||
|
dataset=data_subset,
|
||||||
|
batch_size=batch_size,
|
||||||
|
shuffle=True
|
||||||
|
)
|
||||||
|
|
||||||
|
return data_loader
|
||||||
|
|
||||||
|
|
||||||
|
def load_and_preview_mnist(data_size=100, batch_size=16, seed=42, root='../data'):
|
||||||
|
"""
|
||||||
|
加载并预览 MNIST 数据集。
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
- data_size (int): 加载的数据条目数。
|
||||||
|
- batch_size (int): 每批次处理的数据数量。
|
||||||
|
- seed (int): 随机种子,用于保证数据切割的可重复性。
|
||||||
|
- root (str): 数据集保存的路径。
|
||||||
|
"""
|
||||||
|
train_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root, train=True)
|
||||||
|
|
||||||
|
images, labels = next(iter(train_loader))
|
||||||
|
|
||||||
|
img_grid = torchvision.utils.make_grid(images, nrow=batch_size // 4)
|
||||||
|
img = img_grid.numpy().transpose(1, 2, 0)
|
||||||
|
mean = [0.5]
|
||||||
|
std = [0.5]
|
||||||
|
img = img * std + mean
|
||||||
|
|
||||||
|
print([labels[i] for i in range(batch_size)])
|
||||||
|
plt.imshow(img, cmap='gray')
|
||||||
|
plt.axis('off')
|
||||||
|
plt.show()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
load_and_preview_mnist()
|
||||||
|
|
||||||
|
# Or use custom parameters:
|
||||||
|
# load_and_preview_mnist(data_size=200, batch_size=8, seed=123, root='./data')
|
||||||
@@ -0,0 +1,116 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.optim as optim
|
||||||
|
from torch.utils.data import DataLoader
|
||||||
|
from data_transform import load_mnist, load_and_preview_mnist
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import torchvision.utils as vutils
|
||||||
|
|
||||||
|
|
||||||
|
# 定义卷积神经网络模型
|
||||||
|
class Model(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super(Model, self).__init__()
|
||||||
|
self.conv1 = nn.Sequential(
|
||||||
|
nn.Conv2d(1, 64, kernel_size=3, stride=2, padding=1),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
|
||||||
|
nn.ReLU(),
|
||||||
|
# nn.MaxPool2d(stride=2, kernel_size=2)
|
||||||
|
)
|
||||||
|
|
||||||
|
self.dense = nn.Sequential(
|
||||||
|
nn.Linear(7 * 7 * 128, 512),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Dropout(p=0.8),
|
||||||
|
nn.Linear(512, 10)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.conv1(x) # 卷积处理
|
||||||
|
x = x.view(-1, 7 * 7 * 128) # 对参数实行扁平化处理
|
||||||
|
x = self.dense(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
# 加载和预览数据集
|
||||||
|
def main():
|
||||||
|
# 参数设置
|
||||||
|
data_size = 1000
|
||||||
|
batch_size = 4
|
||||||
|
seed = 42
|
||||||
|
root = '../data'
|
||||||
|
epochs_n = 5
|
||||||
|
learning_rate = 0.001
|
||||||
|
|
||||||
|
# 加载数据加载器
|
||||||
|
train_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root, train=True)
|
||||||
|
test_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root, train=False)
|
||||||
|
|
||||||
|
# 预览数据集(可选)
|
||||||
|
load_and_preview_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root)
|
||||||
|
|
||||||
|
# 初始化模型
|
||||||
|
model = Model()
|
||||||
|
cost = nn.CrossEntropyLoss()
|
||||||
|
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
|
||||||
|
|
||||||
|
# 训练模型
|
||||||
|
for epoch in range(epochs_n):
|
||||||
|
model.train()
|
||||||
|
running_loss = 0.0
|
||||||
|
running_correct = 0
|
||||||
|
print(f"Epoch {epoch + 1}/{epochs_n}")
|
||||||
|
print("-" * 10)
|
||||||
|
|
||||||
|
for data in train_loader:
|
||||||
|
images, labels = data
|
||||||
|
optimizer.zero_grad()
|
||||||
|
outputs = model(images)
|
||||||
|
_, preds = torch.max(outputs.data, 1)
|
||||||
|
loss = cost(outputs, labels)
|
||||||
|
loss.backward()
|
||||||
|
optimizer.step()
|
||||||
|
running_loss += loss.item()
|
||||||
|
running_correct += (preds == labels).sum().item()
|
||||||
|
|
||||||
|
epoch_loss = running_loss / len(train_loader.dataset)
|
||||||
|
epoch_acc = 100 * running_correct / len(train_loader.dataset)
|
||||||
|
|
||||||
|
print(f"Loss: {epoch_loss:.4f}, Train Accuracy: {epoch_acc:.4f}%")
|
||||||
|
|
||||||
|
# 测试模型
|
||||||
|
model.eval()
|
||||||
|
testing_correct = 0
|
||||||
|
with torch.no_grad():
|
||||||
|
for data in test_loader:
|
||||||
|
images, labels = data
|
||||||
|
outputs = model(images)
|
||||||
|
_, preds = torch.max(outputs.data, 1)
|
||||||
|
testing_correct += (preds == labels).sum().item()
|
||||||
|
|
||||||
|
test_acc = 100 * testing_correct / len(test_loader.dataset)
|
||||||
|
print(f"Test Accuracy: {test_acc:.4f}%")
|
||||||
|
|
||||||
|
# 测试数据可视化
|
||||||
|
X_test, y_test = next(iter(test_loader))
|
||||||
|
with torch.no_grad():
|
||||||
|
pred = model(X_test)
|
||||||
|
_, predicted_labels = torch.max(pred, 1)
|
||||||
|
|
||||||
|
print("Predicted Labels:", [i.item() for i in predicted_labels])
|
||||||
|
print("True Labels:", [i.item() for i in y_test])
|
||||||
|
|
||||||
|
# 显示图像
|
||||||
|
img_grid = vutils.make_grid(X_test, nrow=batch_size // 4)
|
||||||
|
img = img_grid.numpy().transpose(1, 2, 0)
|
||||||
|
mean = [0.5]
|
||||||
|
std = [0.5]
|
||||||
|
img = img * std + mean
|
||||||
|
plt.imshow(img, cmap='gray')
|
||||||
|
plt.axis('off')
|
||||||
|
plt.show()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
@@ -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}")
|
||||||
@@ -23,3 +23,35 @@
|
|||||||
### RGB三通道分离并显示。
|
### RGB三通道分离并显示。
|
||||||
|
|
||||||
`Lab-1/src/rgb_img.py`
|
`Lab-1/src/rgb_img.py`
|
||||||
|
|
||||||
|
## **实验二** 基于PyTorch的Minist数据集分类
|
||||||
|
|
||||||
|
### 数据预览测试和数据装载;
|
||||||
|
|
||||||
|
`Lab-2/src/data_load.py`
|
||||||
|
|
||||||
|
### 进行数据处理变换操作;
|
||||||
|
|
||||||
|
`Lab-2/src/data_transform.py`
|
||||||
|
|
||||||
|
### 模型搭建和参数优化;
|
||||||
|
|
||||||
|
`Lab-2/src/model_training.py`
|
||||||
|
|
||||||
|
### 基于pytorch卷积神经网络的模型训练;
|
||||||
|
|
||||||
|
`Lab-2/src/model_training.py`
|
||||||
|
|
||||||
|
### 基于pytorch卷积神经网络的模型测试。
|
||||||
|
|
||||||
|
`Lab-2/src/model_training.py`
|
||||||
|
|
||||||
|
## **实验三** 迁移resnet50猫狗图像分类
|
||||||
|
|
||||||
|
### 通过代码自动下载模型并直接调用;对当前迁移过来的模型进行全连接层的调整;模型训练及测试结果。
|
||||||
|
|
||||||
|
`Lab-3/src/train.py`
|
||||||
|
|
||||||
|
## **实验四** 基于YOLO的目标检测
|
||||||
|
|
||||||
|
## **实验五** 基于YOLOv8的工业缺陷检测实验
|
||||||
|
|||||||
Reference in New Issue
Block a user