上面代码的思路:遍历视频文件,使用Python在cmd中执行ffmpeg命令实现对图片抽取。
新建项目
新建一个图像分类的项目,data里面放数据集,dataset文件夹中自定义数据的读取方法,这次我不采用默认的读取方式,太简单没啥意思。然后再新建train.py和test.py
在项目的根目录新建train.py,然后在里面写训练代码。
导入所需要的库
=======
我这次选用的模型是inception_v3。
import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import SeedlingData
from torch.autograd import Variable
from torchvision.models import inception_v3
设置全局参数
======
设置BatchSize、学习率和epochs,判断是否有cuda环境,如果没有设置为cpu。
设置全局参数
modellr = 1e-4
BATCH_SIZE = 32
EPOCHS = 10
DEVICE = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
图像预处理
=====
在做图像与处理时,train数据集的transform和验证集的transform分开做,train的图像处理出了resize和归一化之外,还可以设置图像的增强,比如旋转、随机擦除等一系列的操作,验证集则不需要做图像增强,另外不要盲目的做增强,不合理的增强手段很可能会带来负作用,甚至出现Loss不收敛的情况。注:inception_v3模型输入的size是3×299×299
transform = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
transform_test = transforms.Compose([
transforms.Resize((299, 299)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
读取数据
====
将数据集解压后放到data文件夹下面,如图:
然后我们在dataset文件夹下面新建 __init__.py和dataset.py,在dataset.py文件夹写入下面的代码:
说一下代码的核心逻辑。
第一步 建立字典,定义类别对应的ID,用数字代替类别。
第二步 在__init__里面编写获取图片路径的方法。测试集只有一层路径直接读取,训练集在train文件夹下面是类别文件夹,先获取到类别,再获取到具体的图