Skip to content

Commit

Permalink
Merge branch 'main' into pr/190
Browse files Browse the repository at this point in the history
  • Loading branch information
lukas-blecher committed Sep 25, 2022
2 parents 80b23d0 + ffcde25 commit 1d06731
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 38 deletions.
4 changes: 2 additions & 2 deletions pix2tex/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,18 +10,18 @@ def main():
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('file', nargs='*', type=str, default=None, help='Predict LaTeX code from image file instead of clipboard (cli 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']:
if arguments.gui or name in ['pix2tex_gui', 'latexocr']:
from .gui import main
else:
from .cli import main
Expand Down
88 changes: 61 additions & 27 deletions pix2tex/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
from PIL import ImageGrab
from PIL import Image
import os
from pathlib import Path
import sys
from typing import Tuple
from typing import List, Optional, Tuple
import atexit
from contextlib import suppress
import logging
Expand Down Expand Up @@ -105,9 +106,9 @@ def __call__(self, img=None, resize=True) -> str:
img = None
if img is None:
if self.last_pic is None:
print('Provide an image.')
return ''
else:
print('\nLast image is: ', end='')
img = self.last_pic.copy()
else:
self.last_pic = img.copy()
Expand Down Expand Up @@ -177,6 +178,36 @@ def output_prediction(pred, args):
webbrowser.open(url)


def predict(model, file, arguments):
img = None
if file:
try:
img = Image.open(os.path.expanduser(file))
except Exception as e:
print(e, end='')
else:
try:
img = ImageGrab.grabclipboard()
except NotImplementedError as e:
print(e, end='')
pred = model(img)
output_prediction(pred, arguments)

def check_file_path(paths:List[Path], wdir:Optional[Path]=None)->List[str]:
files = []
for path in paths:
if type(path)==str:
if path=='':
continue
path=Path(path)
pathsi = ([path] if wdir is None else [path, wdir/path])
for p in pathsi:
if p.exists():
files.append(str(p.resolve()))
elif '*' in path.name:
files.extend([str(pi.resolve()) for pi in p.parent.glob(p.name)])
return list(set(files))

def main(arguments):
path = user_data_dir('pix2tex')
os.makedirs(path, exist_ok=True)
Expand All @@ -187,20 +218,31 @@ def main(arguments):
with suppress(OSError):
readline.read_history_file(history_file)
atexit.register(readline.write_history_file, history_file)
files = check_file_path(arguments.file)
wdir = Path(os.getcwd())
with in_model_path():
model = LatexOCR(arguments)
file = None
if files:
for file in check_file_path(arguments.file, wdir):
print(file + ': ', end='')
predict(model, file, arguments)
model.last_pic = None
with suppress(NameError):
readline.add_history(file)
exit()
pat = re.compile(r't=([\.\d]+)')
while True:
try:
instructions = input('Predict LaTeX code for image ("?"/"h" for help). ')
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()
file = instructions.strip()
ins = file.lower()
t = pat.match(ins)
if ins == 'x':
break
elif ins in ['?', 'h', 'help']:
Expand Down Expand Up @@ -231,26 +273,18 @@ def main(arguments):
setattr(arguments, ins, not getattr(arguments, ins, False))
print('set %s to %s' % (ins, getattr(arguments, ins)))
continue
elif os.path.isfile(os.path.realpath(possible_file)):
file = possible_file
else:
t = re.match(r't=([\.\d]+)', ins)
if t is not None:
t = t.groups()[0]
model.args.temperature = float(t)+1e-8
print('new temperature: T=%.3f' % model.args.temperature)
continue
try:
img = None
if file:
img = Image.open(file)
elif t is not None:
t = t.groups()[0]
model.args.temperature = float(t)+1e-8
print('new temperature: T=%.3f' % model.args.temperature)
continue
files = check_file_path(file.split(' '), wdir)
with suppress(KeyboardInterrupt):
if files:
for file in files:
if len(files)>1:
print(file + ': ', end='')
predict(model, file, arguments)
else:
try:
img = ImageGrab.grabclipboard()
except:
pass
pred = model(img)
output_prediction(pred, arguments)
except KeyboardInterrupt:
pass
predict(model, file, arguments)
file = None
22 changes: 14 additions & 8 deletions pix2tex/gui.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,12 @@ def initUI(self):
self.tempField.setSingleStep(0.1)

# Create snip button
self.snipButton = QPushButton('Snip [Alt+S]', self)
self.snipButton.clicked.connect(self.onClick)
if sys.platform == "darwin":
self.snipButton = QPushButton('Snip [Option+S]', self)
self.snipButton.clicked.connect(self.onClick)
else:
self.snipButton = QPushButton('Snip [Alt+S]', self)
self.snipButton.clicked.connect(self.onClick)

self.shortcut = QShortcut(QKeySequence("Alt+S"), self)
self.shortcut.activated.connect(self.onClick)
Expand Down Expand Up @@ -94,7 +98,10 @@ def toggleProcessing(self, value=None):
text = 'Interrupt'
func = self.interrupt
else:
text = 'Snip [Alt+S]'
if sys.platform == "darwin":
text = 'Snip [Option+S]'
else:
text = 'Snip [Alt+S]'
func = self.onClick
self.retryButton.setEnabled(True)
self.shortcut.setEnabled(not self.isProcessing)
Expand Down Expand Up @@ -276,10 +283,10 @@ def mouseReleaseEvent(self, event):
# account for retina display. #TODO how to check if device is actually using retina display
factor = 2 if sys.platform == "darwin" else 1

x1 = int(min(startPos[0], endPos[0])*factor)
y1 = int(min(startPos[1], endPos[1])*factor)
x2 = int(max(startPos[0], endPos[0])*factor)
y2 = int(max(startPos[1], endPos[1])*factor)
x1 = int(min(startPos[0], endPos[0]))
y1 = int(min(startPos[1], endPos[1]))
x2 = int(max(startPos[0], endPos[0]))
y2 = int(max(startPos[1], endPos[1]))

self.repaint()
QApplication.processEvents()
Expand All @@ -297,7 +304,6 @@ def mouseReleaseEvent(self, event):
self.end = QtCore.QPoint()
self.parent.returnSnip(img)


def main(arguments):
with in_model_path():
if os.name != 'nt':
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

setuptools.setup(
name='pix2tex',
version='0.0.27',
version='0.0.28',
description='pix2tex: Using a ViT to convert images of equations into LaTeX code.',
long_description=long_description,
long_description_content_type='text/markdown',
Expand Down

0 comments on commit 1d06731

Please sign in to comment.