1
G:\DEMO_Project\demo_assist\api-server\
2
├── dataset/
3
│   ├── train/
4
│   │   ├── organic/        # 厨余垃圾图片
5
│   │   └── recyclable/     # 可回收垃圾图片
6
│   └── test/
7
│       ├── organic/
8
│       └── recyclable/
9
├── models/                  # 训练好的模型存这里
10
├── train.py                 # 训练脚本
11
└── ....                     # 其他原有文件

1
import torch
2
import torch.nn as nn
3
import torch.optim as optim
4
from torchvision import datasets, transforms, models
5
from torch.utils.data import DataLoader
6
import os
7
import matplotlib.pyplot as plt
8

9
# ========== 1. 配置参数 ==========
10
DATA_DIR = './dataset'
11
MODEL_SAVE_PATH = './models/waste_classifier.pth'
12
BATCH_SIZE = 32
13
EPOCHS = 10
14
LEARNING_RATE = 0.001
15
IMG_SIZE = 224
16
NUM_CLASSES = 2  # 根据你的类别数修改
17

18
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
19
print(f'使用设备: {device}')
20

21
# ========== 2. 数据预处理 ==========
22
# 训练集：随机裁剪、翻转、旋转，增加数据多样性
23
train_transforms = transforms.Compose([
24
    transforms.RandomResizedCrop(IMG_SIZE),
25
    transforms.RandomHorizontalFlip(),
26
    transforms.RandomRotation(15),
27
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
28
    transforms.ToTensor(),
29
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
30
])
31

32
# 测试集：只需要缩放和归一化
33
test_transforms = transforms.Compose([
34
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
35
    transforms.ToTensor(),
36
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
37
])
38

39
# ========== 3. 加载数据 ==========
40
train_dataset = datasets.ImageFolder(
41
    os.path.join(DATA_DIR, 'train'), transform=train_transforms)
42
test_dataset = datasets.ImageFolder(
43
    os.path.join(DATA_DIR, 'test'), transform=test_transforms)
44

45
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
46
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
47

48
# 自动获取类名映射
49
class_names = train_dataset.classes
50
print(f'类别数量: {len(class_names)}, 类名: {class_names}')
51

52
# ========== 4. 构建模型 ==========
53
# 使用预训练的 ResNet18，只改最后的全连接层
54
model = models.resnet18(weights='DEFAULT')
55
num_features = model.fc.in_features
56
model.fc = nn.Linear(num_features, NUM_CLASSES)
57
model = model.to(device)
58

59
# ========== 5. 损失函数和优化器 ==========
60
criterion = nn.CrossEntropyLoss()
61
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
62

63
# ========== 6. 训练循环 ==========
64
train_losses = []
65
test_accuracies = []
66

67
for epoch in range(EPOCHS):
68
    # --- 训练 ---
69
    model.train()
70
    running_loss = 0.0
71
    for images, labels in train_loader:
72
        images, labels = images.to(device), labels.to(device)
73

74
        optimizer.zero_grad()
75
        outputs = model(images)
76
        loss = criterion(outputs, labels)
77
        loss.backward()
78
        optimizer.step()
79

80
        running_loss += loss.item()
81

82
    avg_loss = running_loss / len(train_loader)
83
    train_losses.append(avg_loss)
84

85
    # --- 测试 ---
86
    model.eval()
87
    correct = 0
88
    total = 0
89
    with torch.no_grad():
90
        for images, labels in test_loader:
91
            images, labels = images.to(device), labels.to(device)
92
            outputs = model(images)
93
            _, predicted = torch.max(outputs, 1)
94
            total += labels.size(0)
95
            correct += (predicted == labels).sum().item()
96

97
    accuracy = 100 * correct / total
98
    test_accuracies.append(accuracy)
99

100
    print(f'Epoch [{epoch+1}/{EPOCHS}], Loss: {avg_loss:.4f}, Test Accuracy: {accuracy:.2f}%')
101

102
# ========== 7. 保存模型 ==========
103
os.makedirs('models', exist_ok=True)
104
torch.save({
105
    'model_state_dict': model.state_dict(),
106
    'class_names': class_names,
107
    'img_size': IMG_SIZE
108
}, MODEL_SAVE_PATH)
109
print(f'模型已保存到 {MODEL_SAVE_PATH}')
110

111
# ========== 8. 绘制训练曲线 ==========
112
plt.figure(figsize=(12, 4))
113
plt.subplot(1, 2, 1)
114
plt.plot(train_losses)
115
plt.title('Training Loss')
116
plt.xlabel('Epoch')
117
plt.subplot(1, 2, 2)
118
plt.plot(test_accuracies)
119
plt.title('Test Accuracy')
120
plt.xlabel('Epoch')
121
plt.ylabel('%')
122
plt.tight_layout()
123
plt.savefig('./models/training_curve.png')
124
plt.show()

1
使用设备: cuda
2
类别数量: 2, 类名: ['organic', 'recyclable']
3
Epoch [1/10], Loss: 0.3456, Test Accuracy: 85.23%
4
...
5
Epoch [10/10], Loss: 0.0213, Test Accuracy: 95.67%
6
模型已保存到 ./models/waste_classifier.pth

1
from flask import Flask, request, jsonify
2
import torch
3
import torch.nn as nn
4
from torchvision import transforms, models
5
from PIL import Image
6
import io
7
import os
8

9
app = Flask(__name__)
10

11
# ========== 1. 加载模型 ==========
12
MODEL_PATH = 'D:/ai-project/models/waste_classifier.pth'
13
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
14

15
# 加载保存的模型信息
16
checkpoint = torch.load(MODEL_PATH, map_location=device)
17
class_names = checkpoint['class_names']
18
img_size = checkpoint.get('img_size', 224)
19

20
# 重建模型结构并加载权重
21
model = models.resnet18(weights=None)
22
num_features = model.fc.in_features
23
model.fc = nn.Linear(num_features, len(class_names))
24
model.load_state_dict(checkpoint['model_state_dict'])
25
model = model.to(device)
26
model.eval()
27

28
print(f'模型已加载，类别: {class_names}, 设备: {device}')
29

30
# ========== 2. 预处理函数 ==========
31
transform = transforms.Compose([
32
    transforms.Resize((img_size, img_size)),
33
    transforms.ToTensor(),
34
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
35
])
36

37
def predict_image(image_bytes):
38
    """输入图片字节流，返回预测类名和置信度"""
39
    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
40
    image_tensor = transform(image).unsqueeze(0).to(device)
41

42
    with torch.no_grad():
43
        outputs = model(image_tensor)
44
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
45
        confidence, predicted_idx = torch.max(probabilities, 1)
46

47
    class_name = class_names[predicted_idx.item()]
48
    confidence_val = confidence.item()
49

50
    return class_name, confidence_val
51

52
# ========== 3. API 接口 ==========
53
@app.route('/health')
54
def health():
55
    return jsonify({
56
        "status": "ok",
57
        "cuda": torch.cuda.is_available(),
58
        "classes": class_names
59
    })
60

61
@app.route('/predict', methods=['POST'])
62
def predict():
63
    # 方案A：接收图片URL（后端下载）
64
    if request.is_json:
65
        data = request.get_json()
66
        image_url = data.get('image_url', '')
67
        if image_url.startswith('http'):
68
            import requests
69
            response = requests.get(image_url, timeout=10)
70
            image_bytes = response.content
71
        else:
72
            # 本地路径
73
            with open('.' + image_url, 'rb') as f:  # 路径前加 . 因为Flask工作目录问题
74
                image_bytes = f.read()
75
        class_name, confidence = predict_image(image_bytes)
76
        return jsonify({
77
            "class": class_name,
78
            "confidence": round(confidence, 4)
79
        })
80

81
    # 方案B：直接上传图片文件
82
    if 'file' in request.files:
83
        file = request.files['file']
84
        image_bytes = file.read()
85
        class_name, confidence = predict_image(image_bytes)
86
        return jsonify({
87
            "class": class_name,
88
            "confidence": round(confidence, 4)
89
        })
90

91
    return jsonify({"error": "no image provided"}), 400
92

93
if __name__ == '__main__':
94
    app.run(host='0.0.0.0', port=5000, debug=False)

1
{
2
  "class": "organic",
3
  "confidence": 0.9578
4
}

1
package com.example.demo.dto;
2

3
public class AiPredictResult {
4
    private String clazz;      // 注意：class 是关键字，用 @JsonProperty
5
    private Double confidence;
6

7
    // Getter / Setter
8
    public String getClazz() { return clazz; }
9
    public void setClazz(String clazz) { this.clazz = clazz; }
10
    public Double getConfidence() { return confidence; }
11
    public void setConfidence(Double confidence) { this.confidence = confidence; }
12
}

1
public AiPredictResult predictByFile(MultipartFile file) throws IOException {
2
    // 把 MultipartFile 转成字节，发给 Flask
3
    // 注意：Spring Boot 在服务器上，笔记本 API 在 localhost:5000（通过 frp）
4
    String url = AI_BASE_URL + "/predict";
5

6
    // 用 RestTemplate 上传文件
7
    HttpHeaders headers = new HttpHeaders();
8
    headers.setContentType(MediaType.MULTIPART_FORM_DATA);
9

10
    // 先存为临时文件
11
    File tempFile = File.createTempFile("upload_", ".jpg");
12
    file.transferTo(tempFile);
13

14
    MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
15
    body.add("file", new FileSystemResource(tempFile));
16

17
    HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);
18
    ResponseEntity<AiPredictResult> response = restTemplate.postForEntity(
19
        url, requestEntity, AiPredictResult.class);
20

21
    // 清理临时文件
22
    tempFile.delete();
23

24
    return response.getBody();
25
}

1
@PostMapping("/upload-and-recognize")
2
public Result<ImageRecord> uploadAndRecognize(@RequestParam("file") MultipartFile file) throws IOException {
3
    // 保存文件（原有逻辑）
4
    Files.createDirectories(Paths.get(uploadDir));
5
    String filename = System.currentTimeMillis() + "_" + file.getOriginalFilename();
6
    File dest = new File(uploadDir + "/" + filename);
7
    file.transferTo(dest);
8

9
    // 调用真实 AI 预测（新逻辑）
10
    AiPredictResult aiResult = aiService.predictByFile(file);
11
    String resultStr = aiResult.getClazz() + " (置信度: " +
12
                       String.format("%.2f", aiResult.getConfidence() * 100) + "%)";
13

14
    // 存数据库
15
    ImageRecord record = new ImageRecord(file.getOriginalFilename(), "/uploads/" + filename);
16
    record.setRecognitionResult(resultStr);
17
    return Result.success(repository.save(record));
18
}