用ResNet18+相移编码搭建一个简单的图像旋转方向识别模型

目录

前言

环境

准备数据集

标注

数据增强

角度编码

加载数据集

搭建模型

训练

验证效果

参考

---

前言

因为工作场景需要，要对图像中的某一行文字，以及这行文字中的上一行文字进行识别。但是文字的方向又是不确定的，如下图（需要识别 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