Merge branch 'lukas-blecher:main' into main

README.md

@@ -24,7 +24,7 @@ There are three ways to get a prediction from an image.
 
 2. Thanks to [@katie-lim](https://github.com/katie-lim), you can use a nice user interface as a quick way to get the model prediction. Just call the GUI with `latexocr`. From here you can take a screenshot and the predicted latex code is rendered using [MathJax](https://www.mathjax.org/) and copied to your clipboard.
 
-    Under linux, it is possible to use the GUI with `gnome-screenshot` which comes with multiple monitor support. You just need to run `latexocr --gnome` (**Note:** you should install `gnome-screenshot` beforehand).
+    Under linux, it is possible to use the GUI with `gnome-screenshot` which comes with multiple monitor support if `gnome-screenshot` was installed beforehand.
 
     ![demo](https://user-images.githubusercontent.com/55287601/117812740-77b7b780-b262-11eb-81f6-fc19766ae2ae.gif)
 

pix2tex/__main__.py

@@ -14,7 +14,6 @@ def main():
     parser.add_argument('-k', '--katex', action='store_true', help='Render the latex code in the browser (cli only)')
 
     parser.add_argument('--gui', action='store_true', help='Use GUI (gui only)')
-    parser.add_argument('--gnome', action='store_true', help='Use gnome-screenshot to capture screenshot (gui only)')
 
     arguments = parser.parse_args()
 

pix2tex/cli.py

@@ -3,13 +3,20 @@
 from PIL import ImageGrab
 from PIL import Image
 import os
+import sys
 from typing import Tuple
+import atexit
+from contextlib import suppress
 import logging
 import yaml
 import re
 
+with suppress(ImportError):
+    import readline
+
 import numpy as np
 import torch
+from torch._appdirs import user_data_dir
 from munch import Munch
 from transformers import PreTrainedTokenizerFast
 from timm.models.resnetv2 import ResNetV2
@@ -133,7 +140,28 @@ def __call__(self, img=None, resize=True) -> str:
 
 
 def output_prediction(pred, args):
-    print(pred, '\n')
+    TERM = os.getenv('TERM', 'xterm')
+    if not sys.stdout.isatty():
+        TERM = 'dumb'
+    try:
+        from pygments import highlight
+        from pygments.lexers import get_lexer_by_name
+        from pygments.formatters import get_formatter_by_name
+
+        if TERM.split('-')[-1] == '256color':
+            formatter_name = 'terminal256'
+        elif TERM != 'dumb':
+            formatter_name = 'terminal'
+        else:
+            formatter_name = None
+        if formatter_name:
+            formatter = get_formatter_by_name(formatter_name)
+            lexer = get_lexer_by_name('tex')
+            print(highlight(pred, lexer, formatter), end='')
+    except ImportError:
+        TERM = 'dumb'
+    if TERM == 'dumb':
+        print(pred)
     if args.show or args.katex:
         try:
             if args.katex:
@@ -150,11 +178,27 @@ def output_prediction(pred, args):
 
 
 def main(arguments):
+    path = user_data_dir('pix2tex')
+    os.makedirs(path, exist_ok=True)
+    history_file = os.path.join(path, 'history.txt')
+    with suppress(NameError):
+        # user can `ln -s /dev/null ~/.local/share/pix2tex/history.txt` to
+        # disable history record
+        with suppress(OSError):
+            readline.read_history_file(history_file)
+        atexit.register(readline.write_history_file, history_file)
     with in_model_path():
         model = LatexOCR(arguments)
         file = None
         while True:
-            instructions = input('Predict LaTeX code for image ("?"/"h" for help). ')
+            try:
+                instructions = input('Predict LaTeX code for image ("?"/"h" for help). ')
+            except KeyboardInterrupt:
+                # TODO: make the last line gray
+                print("")
+                continue
+            except EOFError:
+                break
             possible_file = instructions.strip()
             ins = possible_file.lower()
             if ins == 'x':

pix2tex/gui.py

@@ -1,3 +1,4 @@
+from shutil import which
 import sys
 import os
 import tempfile
@@ -104,7 +105,7 @@ def toggleProcessing(self, value=None):
     @pyqtSlot()
     def onClick(self):
         self.close()
-        if self.args.gnome:
+        if which('gnome-screenshot'):
             self.snip_using_gnome_screenshot()
         else:
             self.snipWidget.snip()

setup.py

@@ -24,6 +24,7 @@
     'torchtext>=0.6.0',
     'imagesize>=1.2.0',
 ]
+highlight = ['pygments']
 
 setuptools.setup(
     name='pix2tex',
@@ -66,10 +67,11 @@
         'albumentations>=0.5.2',
     ],
     extras_require={
-        'all': gui+api+train,
+        'all': gui+api+train+highlight,
         'gui': gui,
         'api': api,
-        'train': train
+        'train': train,
+        'highlight': highlight,
     },
     entry_points={
         'console_scripts': [