目录
前言
环境
准备数据集
标注
数据增强
角度编码
加载数据集
搭建模型
训练
验证效果
参考
因为工作场景需要,要对图像中的某一行文字,以及这行文字中的上一行文字进行识别。但是文字的方向又是不确定的,如下图(需要识别 2213Z 和他的上面一行 301 ),所以萌生了自己搭建一个旋转方向识别模型的想法。
当前虽然已经有不少相对成熟的旋转目标检测开源算法,顺利的话可以顺便把文字识别也做了。但是粗略看了一下,大部分都是只支持检测旋转框,不支持检测旋转目标。即只能给出矩形预测框的旋转角度,不能分辨哪一头是正向,哪一头是反向,需要自己去修改模型。
基于当前各种提取图像特征的CNN网络已经非常成熟,提取文字特征自然应该是不在话下,想来自己搭建一个应该不会特别困难,只要在图像特征上进行角度的回归就可以了。因为场景不会很复杂,所以选择了较为轻量的RestNet18进行试验,最后实现的效果也相当不错。
理论上这个模型可以用于预测任意图像的旋转,用来做旋转验证码摆正应该也是可以的。下面进入正题。
硬件环境
RTX3060 6G 笔记本
软件环境
python 3.10(训练时用)
CUDA 11.6
pytorch 1.12.1
这里我用的是label-studio工具进行标注,然后再进行转格式处理。
GitHub - heartexlabs/label-studio: Label Studio is a multi-type data labeling and annotation tool with standardized output format
按官网的教程安装即可,安装完会启动一个网页服务,在网页上进行标注。
这里建议用conda另外开一个虚拟环境,label-studio支持的python版本不能超过3.9。
conda create -n label python=3.9
conda activate label
# Requires Python >=3.7 <=3.9
pip install label-studio
# Start the server at http://localhost:8080
label-studio
启动后随便注册一下,进入主页面创建一个项目,如下图。
进入项目后,点击右上角的settings,进入项目设置。
选择Labeling Interface标签,设置标注方法。
点击Browse Templates选择标注的模板。
选择里面的Keypoint Labeling,关键点标注即可。
然后创建三个label:l1,l2,v。代表可以表示方向的三个点。
这里的思路是通过点l1和点l2连成一条线L,垂直于L指向点v的方向即代表图像的方向。
因为我的场景是文字的方向,所以点l1和点l2直接标注在首尾两个文字的底部,点v则随便点在文字的上方即可,这样可以确保标注的质量不会太差。
其他场景的话标注两个点代表方向也是可以的。
标注完后效果如图
最后点击export导出,选择导出为JSON-MIN格式。
使用以下方法,将三个点求解成对应的角度,取值范围是-180到180度。
def get_angle(a, b, c):
vector = None
if b[0] == a[0]:
tmpx = c[0] - b[0]
vector = np.array([tmpx, 0])
if b[1] == a[1]:
tmpy = c[1] - b[1]
vector = np.array([0, tmpy])
else:
k_ab = (b[1] - a[1]) / (b[0] - a[0])
b_ab = b[1] - k_ab * b[0]
k_cd = -1 / k_ab
b_cd = c[1] - k_cd * c[0]
x = (b_cd - b_ab) / (k_ab - k_cd)
y = k_ab * x + b_ab
vector = np.array([c[0] - x, c[1] - y])
degree = np.degrees(np.arccos(vector[0] / np.linalg.norm(vector)))
if vector[1] < 0:
degree = -degree
return degree
如果数据集比较少,可以用以下方法手动进行随机旋转,拓展多几份。
# 随机旋转
def enhance(pilimg, angle, num):
imgs = []
assert -180 < angle < 180
for i in range(num):
random_angle = random.randint(-180, 180)
img_copy = pilimg.copy()
angle_copy = angle
img_copy = img_copy.rotate(random_angle)
angle_copy -= float(random_angle)
if angle_copy > 180:
angle_copy = angle_copy - 360.0
elif angle_copy < -180.0:
angle_copy = angle_copy + 360.0
assert -180 < angle_copy < 180
imgs.append((img_copy, angle_copy))
return imgs
也可以在训练的时候,加载数据集后进行随机旋转预处理。
可以另外用imgaug库或其他增强库随机加上干扰。
我训练用的是600张数据,随机干扰加随机旋转,拓展到30000张。
角度编码这里踩过一次坑,原本打算让模型直接预测角度,训练之后总体效果也相当不错,但是到了某个特定角度的时候,总是会出现较大的误差。后来发现是因为训练用的是L1Loss(就是直接相减取绝对值),180和-180度虽然在图像上是一致的,但是损失值却很大。比如170度和-170度,直接相减是相差340,但其实只相差了20度。解决方法来自以下分享
CVPR 2023 有向目标检测角度预测新方法 — 相移编码 | 社区开放麦#45_哔哩哔哩_bilibili
改为用三步相移,将角度用三个cos值来表示,编码和解码的代码实现如下
from math import *
def encode(angle):
theta1 = angle/180*pi
theta2 = theta1 + pi*2/3
theta3 = theta1 + pi*4/3
return cos(theta1), cos(theta2), cos(theta3)
def decode(x1, x2, x3):
num = x1*sin(0) + x2*sin(2*pi/3) + x3*sin(4*pi/3)
dum = x1*cos(0) + x2*cos(2*pi/3) + x3*cos(4*pi/3)
theta = -atan2(num, dum)
return degrees(theta)
最后标注完,将角度编码转换后,标注文件中每一行的格式如下
# 图像名 相移编码x1 相移编码x2 相移编码x3
abc.jpg 0.9976460161233695 -0.5582100427146901 -0.4394359734086792
包装一下dataset
数据集的存放路径为:img子文件夹存放图片,train.txt和test.txt记录训练集和测试集的标注信息。
import torch
from torch.utils.data import Dataset
from PIL import Image
class MyDataSet(Dataset):
def __init__(self, root, transforms=None, anno="train.txt"):
self.root = root
self.img_root = f"{root}/img"
self.anno_path = f"{root}/{anno}"
self.transforms = transforms
with open(self.anno_path, 'r') as f:
self.annos = [line.strip().split() for line in f.readlines()]
def __len__(self):
return len(self.annos)
def __getitem__(self, index):
anno = self.annos[index]
image = Image.open(f"{self.img_root}/{anno[0]}")
target = torch.tensor([float(anno[1]), float(anno[2]), float(anno[3])], dtype=torch.float32)
if self.transforms is not None:
image = self.transforms(image)
return image, target
模型主体部分可以直接用torchvision带的ResNet18,并且可以使用他的预训练权重。
将原本ResNet18最后的全连接层的参数改成(512,3),将从图像中提取的512位特征转换成3位的角度相移编码。
最后根据相移编码的原论文,还要加上sigmoid函数,将每个输出值转换到(-1,1)的范围内(一开始没有加sigmoid,总是训练不到比较好的结果)。
模型代码如下
import torch
import torch.nn as nn
from torchvision import models
class OrientModel(nn.Module):
def __init__(self):
super(OrientModel, self).__init__()
self.model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
self.model.fc = nn.Linear(512, 3)
def forward(self, x):
x = self.model.conv1(x)
x = self.model.bn1(x)
x = self.model.relu(x)
x = self.model.maxpool(x)
x = self.model.layer1(x)
x = self.model.layer2(x)
x = self.model.layer3(x)
x = self.model.layer4(x)
x = self.model.avgpool(x)
x = torch.flatten(x, 1)
x = self.model.fc(x)
x = x.sigmoid() * 2 - 1
return x
import torch
from model import OrientModel
from dataset import MyDataSet
import torch.optim as optim
import torchvision.transforms as transforms
from tqdm import tqdm
import time, os
def train():
save_path = f"weights/{time.strftime('%Y%m%d_%H%M%S')}"
if not os.path.exists(save_path):
os.makedirs(save_path)
# 数据预处理
transform_train = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(0.5, 0.5)
])
transform_test = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(0.5, 0.5)
])
# 加载自定义数据集
trainset = MyDataSet("path/to/traindata", transform_train, "train.txt")
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = MyDataSet("path/to/testdata", transform_test, "test.txt")
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=True, num_workers=2)
net = OrientModel()
criterion = torch.nn.L1Loss()
optimizer = optim.SGD(net.parameters(), lr=0.003, momentum=0.9, weight_decay=5e-5)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
net.train()
num_epochs = 120
for epoch in range(num_epochs):
tbar = tqdm(total=len(trainloader))
tbar.set_description(f'epoch: {epoch + 1}/{num_epochs}')
net.train()
train_loss = 0.0
data_lenth = 0
for i, (inputs, labels) in enumerate(trainloader, 0):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
train_loss += loss.item()
data_lenth += len(inputs)
tbar.set_postfix(loss=f'{loss.item()/len(inputs) * 10000}')
tbar.update(1)
# 打印每轮训练的损失
tbar.set_postfix(loss=f'{train_loss/data_lenth * 10000}')
tbar.close()
if (epoch + 1) % 2 == 0:
# 模型验证
net.eval()
valid_loss = 0.0
data_lenth = 0
with torch.no_grad():
tbar = tqdm(total=len(testloader))
tbar.set_description(f'val:')
for i, (inputs, labels) in enumerate(testloader, 0):
inputs, labels = inputs.to(device), labels.to(device)
outputs = net(inputs)
loss = criterion(outputs, labels)
valid_loss += loss.item()
data_lenth += len(inputs)
tbar.set_postfix(loss=f'{loss.item()/len(inputs) * 10000}')
tbar.update(1)
# 打印每轮验证的准确率
tbar.set_postfix(loss=f'{valid_loss/data_lenth * 10000}')
tbar.close()
torch.save(net.state_dict(), f"{save_path}/epoch_{epoch+1}.pth")
print('weight saved!')
print('Finished training.')
if __name__ == "__main__":
train()
训练了120轮之后,最后在测试集上,预测误差在10度以内的样本达到100%。
让图片一度一度的旋转,看看模型预测出来的是否准确。
import cv2
import torch
from preview_dataset import draw_on_cvimg
import torchvision.transforms as transforms
from model import OrientModel
from PIL import Image
import numpy as np
from angle_coder import decode
model_path = "better_best.pth"
transform = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(0.5, 0.5)
])
net = OrientModel()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_dict = torch.load(model_path, map_location='cpu')
net.load_state_dict(model_dict)
net.to(device)
net.eval()
img = Image.open("path/to/image.jpg")
i = 0
with torch.no_grad():
while True:
pilimg = img.rotate(i)
input = transform(pilimg).unsqueeze(0).to(device)
outputs = net(input)
angle = decode(*outputs[0])
cvimg = cv2.cvtColor(np.asarray(pilimg), cv2.COLOR_RGB2BGR)
cvimg = draw_on_cvimg(cvimg, angle)
cv2.imshow('a', cvimg)
cv2.waitKey(2)
i += 1
效果如下
Phase-Shifting Coder: Predicting Accurate Orientation in Oriented Object Detection