大家好,我是一个普通一本的大一计科生。这篇文章不是大佬的教学,而是一个小白踩坑一个月的真实记录。一个月前我连 import torch 都会报错,现在做出了一个能上传图片判断猫狗的网站。

 如果你也刚入坑AI,想做一个完整的项目但不知道从哪下手,这篇文章应该能帮到你。

 

项目效果

 一、项目整体思路

这个项目分三步:数据 → 模型 → 网站

爬虫下载猫狗图片 → 训练ResNet18模型 → Flask搭建网站 → 上传图片识别

技术栈就是PyTorch加Flask,没有用很花哨的东西。

技术栈一览

爬虫 requests + re 轻量,不用Scrapy那么重

训练 PyTorch + ResNet18 预训练模型,省时间

网站 Flask + HTML 简单,几行代码就能跑

部署 本地运行 先跑通再说

二、第一步:爬虫采集数据

为啥不直接用现成的数据集?

网上有猫狗数据集,但我想自己控制数据来源。而且爬虫写好了,以后想加"兔子""鸟"之类的,改个关键词就能跑。

代码实现

python

import requests

import re

import os

import time

 

def download_baidu_images(keyword, save_dir, num=30):

    os.makedirs(save_dir, exist_ok=True)

    

    headers = {

        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'

    }

    

    downloaded = 0

    page = 0

    

    while downloaded < num:

        url = f'https://image.baidu.com/search/flip?tn=baiduimage&word={keyword}&pn={page*20}'

        resp = requests.get(url, headers=headers)

        

        # 从返回的html里抠图片链接

        img_urls = re.findall(r'"objURL":"(.*?)"', resp.text)

        

        for img_url in img_urls:

            if downloaded >= num:

                break

            try:

                img_data = requests.get(img_url, timeout=10)

                with open(f'{save_dir}/{keyword}_{downloaded+1}.jpg', 'wb') as f:

                    f.write(img_data.content)

                downloaded += 1

                print(f'已下载 {downloaded}/{num}')

                time.sleep(0.5) # 别太快,会被封

            except:

                continue

        page += 1

        time.sleep(1)

 

download_baidu_images('猫', './dataset/猫', 30)

download_baidu_images('狗', './dataset/狗', 30)

 

踩过的坑

 

1. 百度会封IP:加 time.sleep 和 User-Agent 会好很多。封了等几分钟再试。

2. 有些链接不是图片:下载下来打不开的,直接 except 跳过就行。

3. 图片格式不对:统一保存成 .jpg,PIL都能打开。

三、第二步:训练模型

为啥选ResNet18?

ResNet18够用,ResNet50我的电脑带不动。而且PyTorch一行代码就能下载预训练好的版本,直接改最后一层就行。

数据预处理

python

from torchvision import transforms

transform = transforms.Compose([

    transforms.Resize((224, 224)), # ResNet输入要求

    transforms.RandomHorizontalFlip(), # 随机翻转,相当于白嫖数据

    transforms.ToTensor(),

    transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225]) # ImageNet标准

])

 

 

模型定义

 

python

import torch.nn as nn

import torchvision.models as models

 

model = models.resnet18(pretrained=True) # 下载预训练权重

model.fc = nn.Linear(512, 2) # 改成2分类:猫和狗

 

 

训练代码

 

python

import torch.optim as optim

 

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)

 

for epoch in range(10):

    model.train()

    total_loss = 0

    for images, labels in train_loader:

        optimizer.zero_grad()

        outputs = model(images)

        loss = criterion(outputs, labels)

        loss.backward()

        optimizer.step()

        total_loss += loss.item()

    

    print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')

训练结果

 

我只有60张图(猫30,狗30),训练了10个epoch:

 

· 训练准确率:100%(已经背下来了)

· 验证准确率:83.33%(12张验证集猜对了10张)

 

数据量太小,有过拟合。如果你有200张以上,准确率应该能到90%以上。

 

四、第三步:部署成网站

 

为啥用Flask?

 

因为简单。整个后端就几十行代码,不需要学Django那种大家伙。

 

后端代码

 

python

from flask import Flask, request, render_template_string

import torch

from PIL import Image

 

app = Flask(__name__)

model = load_model() # 启动时加载一次就行

 

@app.route('/')

def index():

    return render_template_string(HTML) # 前端代码直接写在里面

 

@app.route('/predict', methods=['POST'])

def predict():

    file = request.files['file']

    img = Image.open(file.stream).convert('RGB')

    # 预处理 → 预测 → 返回结果

    return {'class': '猫', 'confidence': 0.99}

```

 

前端界面

 

就是一个HTML页面:上传框 + 按钮 + 结果显示区。支持点击上传和拖拽上传。

 

前端代码比较长,核心就是用 FormData 发请求:

 

```javascript

fetch('/predict', {

    method: 'POST',

    body: formData

})

.then(res => res.json())

.then(data => {

    document.getElementById('result').innerHTML = 

        `识别结果:${data.class},置信度:${(data.confidence*100).toFixed(2)}%`;

});

```

 

启动

 

```bash

python app.py

```

 

浏览器打开 http://127.0.0.1:5000,上传一张图,就能看到结果了。

五、项目优化方向

 

目前只是个demo,想做得更好可以从这几个方向下手:

 

1. 增加数据量:爬虫多跑几次,每类100张以上,准确率能到90%以上

2. 加更多类别:爬虫关键词改成"鸟""兔子",重新训练就行

3. 部署到云服务器:让别人可以点链接直接访问,不用自己跑代码

六、一些真心话

 

一个月能做成这样?

 

能。前提是每天坚持2-3小时。我大一课不多,晚上和周末都在写代码。

 

最开始也看不懂 view(-1, 784) 是啥意思,squeeze() 为啥要去掉维度1。硬着头皮跑通,跑多了就懂了。

 

给同样在入门的同学几点建议

 

1. 不要一直看课:看十节课不如跑通一个项目

2. 代码一定要手敲:复制粘贴看起来快,但坑都踩不到,等于白学

3. 遇到报错读英文:别直接复制到百度,试着看最后一行错误信息,慢慢就能看懂

4. 做完一定要写文章:写下来才是自己的。我在写这篇文章的时候,又把代码理了一遍,发现好几个之前没注意的点

 

关于大一找实习

 

有人问我现在能不能去实习。实话实说:挺难的。不是因为能力不够,是企业根本不招大一。

 

但这不代表暑假没事干。我计划暑假:刷LeetCode + 做第二个项目(试试YOLO目标检测)+ 写技术博客。大二暑假再去找实习。

 

附:环境配置

# 创建虚拟环境

conda create -n ai_env python=3.9

conda activate ai_env

 

# 安装依赖

pip install torch torchvision flask pillow requests tqdm

 

# 运行

python spider.py --keyword "猫" --num 30 # 爬数据

python train.py # 训练

python app.py # 启动网站

`

 

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐