1
项目文件夹/
2
  train.py
3
  predict.py
4
  garbage-dataset/
5
    battery/
6
    biological/
7
    cardboard/
8
    clothes/
9
    glass/
10
    metal/
11
    paper/
12
    plastic/
13
    shoes/
14
    trash/

1
import torch
2
import torch.nn as nn
3
from torchvision import datasets, transforms, models
4
from torch.utils.data import DataLoader, Subset
5
from sklearn.metrics import classification_report
6

7
# 数据预处理
8
train_transform = transforms.Compose([
9
    transforms.Resize((128, 128)),
10
    transforms.RandomHorizontalFlip(),
11
    transforms.RandomRotation(15),
12
    transforms.ToTensor(),
13
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
14
])
15

16
val_transform = transforms.Compose([
17
    transforms.Resize((128, 128)),
18
    transforms.ToTensor(),
19
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
20
])
21

22
# 加载数据集（10个类别，共19762张图片，8:2划分）
23
data_dir = "./garbage-dataset"
24
train_data = datasets.ImageFolder(data_dir, transform=train_transform)
25
val_data = datasets.ImageFolder(data_dir, transform=val_transform)
26
class_names = train_data.classes
27

28
n_total = len(train_data)
29
n_train = int(n_total * 0.8)
30
n_val = n_total - n_train
31
indices = torch.randperm(n_total).tolist()
32
train_set = Subset(train_data, indices[:n_train])
33
val_set = Subset(val_data, indices[n_train:])
34
train_loader = DataLoader(train_set, batch_size=16, shuffle=True)
35
val_loader = DataLoader(val_set, batch_size=16, shuffle=False)
36

37
# 加载预训练模型并修改分类头
38
model = models.mobilenet_v2(weights='IMAGENET1K_V1')
39
for param in model.features.parameters():
40
    param.requires_grad = False
41
model.classifier[1] = nn.Linear(model.last_channel, len(class_names))
42

43
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
44
model = model.to(device)
45

46
# 训练配置
47
criterion = nn.CrossEntropyLoss()
48
optimizer = torch.optim.Adam(model.classifier.parameters(), lr=0.001)
49
EPOCHS = 10
50
best_acc = 0.0
51

52
for epoch in range(EPOCHS):
53
    # 训练阶段
54
    model.train()
55
    total_loss = 0
56
    correct = 0
57
    for imgs, labels in train_loader:
58
        imgs, labels = imgs.to(device), labels.to(device)
59
        optimizer.zero_grad()
60
        outputs = model(imgs)
61
        loss = criterion(outputs, labels)
62
        loss.backward()
63
        optimizer.step()
64
        total_loss += loss.item()
65
        correct += (outputs.argmax(1) == labels).sum().item()
66
    train_acc = correct / n_train
67

68
    # 验证阶段
69
    model.eval()
70
    val_correct = 0
71
    with torch.no_grad():
72
        for imgs, labels in val_loader:
73
            imgs, labels = imgs.to(device), labels.to(device)
74
            outputs = model(imgs)
75
            val_correct += (outputs.argmax(1) == labels).sum().item()
76
    val_acc = val_correct / n_val
77

78
    print(f"Epoch {epoch+1}/{EPOCHS} | Loss: {total_loss/len(train_loader):.3f} | Train Acc: {train_acc:.3f} | Val Acc: {val_acc:.3f}")
79
    if val_acc > best_acc:
80
        best_acc = val_acc
81
        torch.save(model.state_dict(), 'best_model.pth')
82

83
print(f"最佳验证准确率: {best_acc:.3f}")

1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
2
model = model.to(device)

1
device = torch.device('cpu')
2
model = model.to(device)

1
import torch
2
import torch.nn as nn
3
from torchvision import transforms, models
4
from PIL import Image
5
import matplotlib.pyplot as plt
6

7
class_names = ['battery', 'biological', 'cardboard', 'clothes', 'glass', 'metal', 'paper', 'plastic', 'shoes', 'trash']
8

9
model = models.mobilenet_v2(weights=None)
10
model.classifier[1] = nn.Linear(model.last_channel, len(class_names))
11
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
12
model.eval()
13

14
transform = transforms.Compose([
15
    transforms.Resize((128, 128)),
16
    transforms.ToTensor(),
17
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
18
])
19

20
img_path = 'test.jpg'
21
img = Image.open(img_path).convert('RGB')
22
tensor = transform(img).unsqueeze(0)
23

24
with torch.no_grad():
25
    outputs = model(tensor)
26
    probs = torch.softmax(outputs, dim=1)[0]
27
    pred = probs.argmax().item()
28

29
print(f"预测结果: {class_names[pred]}")
30
print(f"置信度: {probs[pred]:.2%}")
31

32
plt.figure(figsize=(5,5))
33
plt.imshow(img)
34
plt.axis('off')
35
plt.title(f"预测: {class_names[pred]}\n置信度: {probs[pred]:.2%}")
36
plt.savefig('result.jpg')
37
plt.show()