爬虫部分基于 PixivCrawler 修改实现, 涩图识别分类部分使用 ConvNeXt 作为backbone的分类模型实现, 性能优于Trasnformer类模型。
环境配置参考 ConvNeXt
需要 pytorch==1.8 timm==0.3.2
下载miniconda,创建新python环境并激活
conda create -n pixivai python=3.9
conda activate pixivai
安装pytorch
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch-lts -c conda-forge
# 没有N卡的用这个
conda install pytorch torchvision torchaudio cpuonly -c pytorch-lts
安装其他依赖
pip install -r requirements.txt
下载预训练权重放在ckpt/
文件夹内:
根据 PixivCrawler 的说明配置爬虫,设置账号和cookie,设置要爬的内容。
pixiv_crawler/config.py
中配置爬虫基本参数。
运行命令启动AI爬虫:
# 不加关键字默认爬日榜
python AIcrawler.py --ckpt 模型权重 --n_images 总图像个数 [--keyword 关键字]
准备至少5000张图。
用labeler.py
打标签,数据集标签会储存为json格式。
或
把不同类别放入不同文件夹,用labeler_folder.py
一键打标签。
images
|--0
| |--1.png
| |--2.png
|
|--1
用data_proc.py
划分训练集和测试集,并对图像进行预处理。
修改参数,运行脚本训练:
python train.sh
训练参数设置参考 ConvNeXt