tps_stn_pytorch

tps_stn_pytorch

Implémentation Pytorch du réseau de transformateurs spatiales (STN) avec une spline à plaques minces (TPS).

Introduction

STN est une puissante architecture de réseau neuronal proposée par DeepMind dans [1]. STN obtient une véritable invariance spatiale en rectifiant automatiquement les images d'entrée avant qu'elles ne soient alimentées dans un réseau de classification normal. La partie la plus étonnante de STN est qu'il est un différentiel de bout en bout et peut être directement branché sur les architectures de réseau existantes (Alexnet, Resnet, etc.), sans aucune supervision supplémentaire.

Expériences de papier STN d'origine [1] sur trois formes de transformation spécifiques: transformation affine, transformation projective et transformation de spline à plaque mince (TPS) . Parmi eux, je pense que TPS est la traduction la plus puissante car elle peut déformer une image de manière arbitraire. Comme indiqué ci-dessous, je peux déformer mon avatar

dans

TPS-STN a été utilisé dans l'application OCR [2]. Dans cet article, TPS-STN consiste à rectifier automatiquement les images de texte déformées, avant qu'elles ne soient alimentées dans un modèle de reconnaissance de texte OCR normal:

Dépendances

• Python3
• Pytorch
• Torchion
• Nombant
• Matplotlib
• Oreiller / pilon
• image

J'utilise imageio pour créer une visualisation GIF. Installez-le simplement par pip install imageio .

Courir

 python mnist_train.py --model unbounded_stn --angle 90 --grid_size 4
python mnist_visualize.py --model unbounded_stn --angle 90 --grid_size 4
python mnist_make_gif.py --model unbounded_stn --angle 90 --grid_size 4

Ensuite, Png et GIF ResUtls seront enregistrés dans ./image/unbounded_stn_angle60_grid4/ et ./gif/unbounded_stn_angle60_grid4/ .

Vous pouvez essayer d'autres combinaisons d'architecture du modèle, un angle de rotation aléatoire MNIST et une taille de grille TPS. Détails ci-dessous.

Arguments

Il y a trois arguments contrôlables: --model , --angle , --grid_size .

--model : str, requis

• Avec no_stn , le module STN est jeté et un seul classificateur CNN demeure.
• Avec bounded_stn , la sortie du réseau de localisation est pressée à [-1, 1] par F.tanh , comme cela a été fait dans [2]
• Avec unbounded_stn , la sortie du réseau de locolizaition n'est pas pressée

--angle : int, par défaut = 60

• Les échantillons dans l'ensemble de données MNIST seront tournés par des angles aléatoires dans [-angle, angle]

--grid_size : int, default = 4

• Utiliser (grid_size x grid_size) points de contrôle pour définir la transformation de la spline à plaque mince

Courbes de précision de test

Visualiser TPS-STN

Les résultats avec angle = 90 sont généralement mauvais:

Les résultats avec bounded_stn sont mauvais si grid_size <= 3 :

Mais ok si grid_size >= 4 :

Les résultats avec unbounded_stn sont OK:

Bien sûr, il y a toujours de mauvais cas dans chaque combinaison. Vous pouvez télécharger tous mes GIF à partir de Baidu NetDisk (File Size 2G).

Référence

[1] Réseaux de transformateurs spatiaux

[2] Reconnaissance robuste de texte de scène avec recti fi cation automatique

[3] 数值方法 —— 薄板样条插值 （Spline à plaque mince）