用 MobileNet 教 AI 分垃圾？这个项目很适合练手也很有用！

import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision.models import mobilenet_v2

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

dataset = ImageFolder(root='data/train', transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 模型准备
model = mobilenet_v2(pretrained=True)
model.classifier[1] = nn.Linear(model.last_channel, 4)

# 训练配置
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

# 简单训练过程
for epoch in range(5):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

model.eval()
img = Image.open("test.jpg").convert("RGB")
input_tensor = transform(img).unsqueeze(0)
output = model(input_tensor)
pred = torch.argmax(output, dim=1)

from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io

app = FastAPI()

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    input_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        output = model(input_tensor)
    pred = torch.argmax(output, dim=1).item()
    return {"class": label_map[pred]}

<input type="file" id="fileInput" />
<button onclick="upload()">识别垃圾类别</button>

<script>
async function upload() {
  const file = document.getElementById('fileInput').files[0]
  const formData = new FormData()
  formData.append('file', file)

  const res = await fetch('/predict/', {
    method: 'POST',
    body: formData
  })
  const result = await res.json()
  alert("识别结果：" + result.class)
}
</script>