Merge pull request lukas-blecher#179 from Freed-Wu/main

use same parser to parse cli and gui arguments

Author: lukas-blecher

pix2tex/__main__.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+def main():
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument('-t', '--temperature', type=float, default=.333, help='Softmax sampling frequency')
+    parser.add_argument('-c', '--config', type=str, default='settings/config.yaml', help='path to config file')
+    parser.add_argument('-m', '--checkpoint', type=str, default='checkpoints/weights.pth', help='path to weights file')
+    parser.add_argument('--no-cuda', action='store_true', help='Compute on CPU')
+    parser.add_argument('--no-resize', action='store_true', help='Resize the image beforehand')
+
+    parser.add_argument('-s', '--show', action='store_true', help='Show the rendered predicted latex code (cli only)')
+    parser.add_argument('-f', '--file', type=str, default=None, help='Predict LaTeX code from image file instead of clipboard (cli only)')
+    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()
+
+    import os
+    import sys
+
+    name = os.path.split(sys.argv[0])[-1]
+    if arguments.gui or arguments.gnome or name in ['pix2tex_gui', 'latexocr']:
+        from .gui import main
+    else:
+        from .cli import main
+    main(arguments)
+
+
+if __name__ == '__main__':
+    main()

pix2tex/cli.py

@@ -4,7 +4,6 @@
 from PIL import Image
 import os
 from typing import Tuple
-import argparse
 import logging
 import yaml
 import re
@@ -150,17 +149,7 @@ def output_prediction(pred, args):
             webbrowser.open(url)
 
 
-def main():
-    parser = argparse.ArgumentParser(description='Use model')
-    parser.add_argument('-t', '--temperature', type=float, default=.333, help='Softmax sampling frequency')
-    parser.add_argument('-c', '--config', type=str, default='settings/config.yaml')
-    parser.add_argument('-m', '--checkpoint', type=str, default='checkpoints/weights.pth')
-    parser.add_argument('-s', '--show', action='store_true', help='Show the rendered predicted latex code')
-    parser.add_argument('-f', '--file', type=str, default=None, help='Predict LaTeX code from image file instead of clipboard')
-    parser.add_argument('-k', '--katex', action='store_true', help='Render the latex code in the browser')
-    parser.add_argument('--no-cuda', action='store_true', help='Compute on CPU')
-    parser.add_argument('--no-resize', action='store_true', help='Resize the image beforehand')
-    arguments = parser.parse_args()
+def main(arguments):
     with in_model_path():
         model = LatexOCR(arguments)
         file = None
@@ -221,7 +210,3 @@ def main():
             except KeyboardInterrupt:
                 pass
             file = None
-
-
-if __name__ == "__main__":
-    main()

pix2tex/gui.py

@@ -1,6 +1,5 @@
 import sys
 import os
-import argparse
 import tempfile
 from PyQt5 import QtCore, QtGui
 from PyQt5.QtCore import QObject, Qt, pyqtSlot, pyqtSignal, QThread
@@ -298,20 +297,8 @@ def mouseReleaseEvent(self, event):
         self.parent.returnSnip(img)
 
 
-def main():
-    parser = argparse.ArgumentParser(description='GUI arguments')
-    parser.add_argument('-t', '--temperature', type=float, default=.2, help='Softmax sampling frequency')
-    parser.add_argument('-c', '--config', type=str, default='settings/config.yaml', help='path to config file')
-    parser.add_argument('-m', '--checkpoint', type=str, default='checkpoints/weights.pth', help='path to weights file')
-    parser.add_argument('--no-cuda', action='store_true', help='Compute on CPU')
-    parser.add_argument('--no-resize', action='store_true', help='Resize the image beforehand')
-    parser.add_argument('--gnome', action='store_true', help='Use gnome-screenshot to capture screenshot')
-    arguments = parser.parse_args()
+def main(arguments):
     with in_model_path():
         app = QApplication(sys.argv)
         ex = App(arguments)
         sys.exit(app.exec_())
-
-
-if __name__ == '__main__':
-    main()

setup.py

@@ -73,10 +73,10 @@
     },
     entry_points={
         'console_scripts': [
-            'pix2tex_gui = pix2tex.gui:main',
-            'pix2tex_cli = pix2tex.cli:main',
-            'latexocr = pix2tex.gui:main',
-            'pix2tex = pix2tex.cli:main',
+            'pix2tex_gui = pix2tex.__main__:main',
+            'pix2tex_cli = pix2tex.__main__:main',
+            'latexocr = pix2tex.__main__:main',
+            'pix2tex = pix2tex.__main__:main',
         ],
     },
     classifiers=[