+extra train tag, mv transformations to sep. file

Author: lukas-blecher

README.md

@@ -40,6 +40,7 @@ Always double check the result carefully. You can try to redo the prediction wit
 
 ## Training the model [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lukas-blecher/LaTeX-OCR/blob/main/notebooks/LaTeX_OCR_training.ipynb)
 
+Install a couple of dependencies `pip install pix2tex[train]`.
 1. First we need to combine the images with their ground truth labels. I wrote a dataset class (which needs further improving) that saves the relative paths to the images with the LaTeX code they were rendered with. To generate the dataset pickle file run 
 
 ```

notebooks/LaTeX_OCR_training.ipynb

@@ -35,7 +35,7 @@
         "id": "r396ah-Q3EQc"
       },
       "source": [
-        "!pip install pix2tex -qq"
+        "!pip install pix2tex[train] -qq"
       ],
       "execution_count": null,
       "outputs": []

pix2tex/cli.py

@@ -1,4 +1,4 @@
-from pix2tex.dataset.dataset import test_transform
+from pix2tex.dataset.transforms import test_transform
 import pandas.io.clipboard as clipboard
 from PIL import ImageGrab
 from PIL import Image

pix2tex/dataset/dataset.py

@@ -1,6 +1,3 @@
-from tempfile import tempdir
-import albumentations as alb
-from albumentations.pytorch import ToTensorV2
 import torch
 import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
@@ -17,33 +14,8 @@
 from tqdm.auto import tqdm
 
 from pix2tex.utils.utils import in_model_path
+from pix2tex.dataset.transforms import train_transform, test_transform
 
-train_transform = alb.Compose(
-    [
-        alb.Compose(
-            [alb.ShiftScaleRotate(shift_limit=0, scale_limit=(-.15, 0), rotate_limit=1, border_mode=0, interpolation=3,
-                                  value=[255, 255, 255], p=1),
-             alb.GridDistortion(distort_limit=0.1, border_mode=0, interpolation=3, value=[255, 255, 255], p=.5)], p=.15),
-        # alb.InvertImg(p=.15),
-        alb.RGBShift(r_shift_limit=15, g_shift_limit=15,
-                     b_shift_limit=15, p=0.3),
-        alb.GaussNoise(10, p=.2),
-        alb.RandomBrightnessContrast(.05, (-.2, 0), True, p=0.2),
-        alb.ImageCompression(95, p=.3),
-        alb.ToGray(always_apply=True),
-        alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
-        # alb.Sharpen()
-        ToTensorV2(),
-    ]
-)
-test_transform = alb.Compose(
-    [
-        alb.ToGray(always_apply=True),
-        alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
-        # alb.Sharpen()
-        ToTensorV2(),
-    ]
-)
 
 
 class Im2LatexDataset:

pix2tex/dataset/transforms.py

@@ -0,0 +1,29 @@
+import albumentations as alb
+from albumentations.pytorch import ToTensorV2
+
+train_transform = alb.Compose(
+    [
+        alb.Compose(
+            [alb.ShiftScaleRotate(shift_limit=0, scale_limit=(-.15, 0), rotate_limit=1, border_mode=0, interpolation=3,
+                                  value=[255, 255, 255], p=1),
+             alb.GridDistortion(distort_limit=0.1, border_mode=0, interpolation=3, value=[255, 255, 255], p=.5)], p=.15),
+        # alb.InvertImg(p=.15),
+        alb.RGBShift(r_shift_limit=15, g_shift_limit=15,
+                     b_shift_limit=15, p=0.3),
+        alb.GaussNoise(10, p=.2),
+        alb.RandomBrightnessContrast(.05, (-.2, 0), True, p=0.2),
+        alb.ImageCompression(95, p=.3),
+        alb.ToGray(always_apply=True),
+        alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
+        # alb.Sharpen()
+        ToTensorV2(),
+    ]
+)
+test_transform = alb.Compose(
+    [
+        alb.ToGray(always_apply=True),
+        alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
+        # alb.Sharpen()
+        ToTensorV2(),
+    ]
+)

setup.py

@@ -19,10 +19,15 @@
     'uvicorn[standard]',
     'python-multipart'
 ]
+train = [
+    'python-Levenshtein>=0.12.2',
+    'torchtext>=0.6.0',
+    'imagesize>=1.2.0',
+]
 
 setuptools.setup(
     name='pix2tex',
-    version='0.0.22',
+    version='0.0.23',
     description='pix2tex: Using a ViT to convert images of equations into LaTeX code.',
     long_description=long_description,
     long_description_content_type='text/markdown',
@@ -58,15 +63,13 @@
         'PyYAML>=5.4.1',
         'pandas>=1.0.0',
         'timm',
-        'python-Levenshtein>=0.12.2',
-        'torchtext>=0.6.0',
         'albumentations>=0.5.2',
-        'imagesize>=1.2.0',
     ],
     extras_require={
-        'all': gui+api,
+        'all': gui+api+train,
         'gui': gui,
-        'api': api
+        'api': api,
+        'train': train
     },
     entry_points={
         'console_scripts': [