diff --git a/pix2tex/api/streamlit.py b/pix2tex/api/streamlit.py index 128d55a..1342aec 100644 --- a/pix2tex/api/streamlit.py +++ b/pix2tex/api/streamlit.py @@ -1,32 +1,79 @@ import requests from PIL import Image -import streamlit +import streamlit as st +from st_img_pastebutton import paste +from io import BytesIO +import base64 -if __name__ == '__main__': - streamlit.set_page_config(page_title='LaTeX-OCR') - streamlit.title('LaTeX OCR') - streamlit.markdown('Convert images of equations to corresponding LaTeX code.\n\nThis is based on the `pix2tex` module. Check it out [![github](https://img.shields.io/badge/LaTeX--OCR-visit-a?style=social&logo=github)](https://github.com/lukas-blecher/LaTeX-OCR)') - uploaded_file = streamlit.file_uploader( - 'Upload an image an equation', - type=['png', 'jpg'], +def encode_image(file): + _, encoded = file.split(",", 1) + binary_data = base64.b64decode(encoded) + bytes_data = BytesIO(binary_data) + return bytes_data + + +if __name__ == "__main__": + st.set_page_config(page_title="LaTeX-OCR") + st.title("LaTeX OCR") + st.markdown( + "Convert images of equations to corresponding LaTeX code.\n\nThis is based on the `pix2tex` module. Check it out [![github](https://img.shields.io/badge/LaTeX--OCR-visit-a?style=social&logo=github)](https://github.com/lukas-blecher/LaTeX-OCR)" ) - if uploaded_file is not None: - image = Image.open(uploaded_file) - streamlit.image(image) - else: - streamlit.text('\n') + source = st.radio( + "Choose the source of the image", + options=["Paste", "Upload"], + ) + + image = None + + if source == "Upload": + uploaded_file = st.file_uploader( + "Upload an image of an equation", + type=["png", "jpg"], + ) - if streamlit.button('Convert'): - if uploaded_file is not None and image is not None: - with streamlit.spinner('Computing'): - response = requests.post('http://127.0.0.1:8502/predict/', files={'file': uploaded_file.getvalue()}) + if uploaded_file is not None: + st.image(Image.open(uploaded_file)) + image = uploaded_file.getvalue() + + if source == "Paste": + pasted_file = paste("Paste an image of an equation") + + if pasted_file is not None: + image = encode_image(pasted_file) + st.image(image) + + if st.button("Convert"): + if image is not None: + with st.spinner("Computing"): + response = requests.post( + "http://127.0.0.1:8502/predict/", files={"file": image} + ) if response.ok: - latex_code = response.json() - streamlit.code(latex_code, language='latex') - streamlit.markdown(f'$\\displaystyle {latex_code}$') - else: - streamlit.error(response.text) - else: - streamlit.error('Please upload an image.') + st.session_state.latex_code = response.json() + + if 'latex_code' in st.session_state: + + st.subheader("Preview") + st.markdown(f"$\\displaystyle {st.session_state.latex_code}$") + + format_option = st.selectbox( + "Select format to copy:", + options=[ + "Raw LaTeX", + "Inline Math ($....$)", + "Display Math ($$....$$)" + ] + ) + + formatted_code = st.session_state.latex_code + if format_option == "Inline Math ($....$)": + formatted_code = f"${st.session_state.latex_code}$" + elif format_option == "Display Math ($$....$$)": + formatted_code = f"$${st.session_state.latex_code}$$" + + st.code(formatted_code, language="latex") + + elif image is None and st.session_state.get('_button_clicked', False): + st.error("No image selected") \ No newline at end of file diff --git a/pix2tex/resources/resources.py b/pix2tex/resources/resources.py index ddab65f..7d383f0 100644 --- a/pix2tex/resources/resources.py +++ b/pix2tex/resources/resources.py @@ -3,7 +3,7 @@ # Created by: The Resource Compiler for Qt version 6.4.2 # WARNING! All changes made in this file will be lost! -from PySide6 import QtCore +from PyQt6 import QtCore qt_resource_data = b"\ \x00\x02*\x89\ diff --git a/pix2tex/setup.py b/pix2tex/setup.py new file mode 100644 index 0000000..34c22d7 --- /dev/null +++ b/pix2tex/setup.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python + +import setuptools + +# read the contents of your README file +from pathlib import Path +this_directory = Path(__file__).parent +long_description = (this_directory / 'README.md').read_text(encoding='utf-8') + +gui = [ + 'PyQt6', + 'PyQt6-WebEngine', + 'pyside6', + 'pynput', + 'screeninfo', +] +api = [ + 'streamlit>=1.8.1', + 'fastapi>=0.75.2', + 'uvicorn[standard]', + 'python-multipart', + 'st_img_pastebutton>=0.0.3', +] +train = [ + 'python-Levenshtein>=0.12.2', + 'torchtext>=0.6.0', + 'imagesize>=1.2.0', +] +highlight = ['pygments'] + +setuptools.setup( + name='pix2tex', + version='0.1.2', + description='pix2tex: Using a ViT to convert images of equations into LaTeX code.', + long_description=long_description, + long_description_content_type='text/markdown', + author='Lukas Blecher', + author_email='luk.blecher@gmail.com', + url='https://github.com/lukas-blecher/LaTeX-OCR/', + license='MIT', + keywords=[ + 'artificial intelligence', + 'deep learning', + 'image to text' + ], + packages=setuptools.find_packages(), + package_data={ + 'pix2tex': [ + 'resources/*', + 'model/settings/*.yaml', + 'model/dataset/*.json', + ] + }, + install_requires=[ + 'tqdm>=4.47.0', + 'munch>=2.5.0', + 'torch>=1.7.1', + 'opencv_python_headless>=4.1.1.26', + 'requests>=2.22.0', + 'einops>=0.3.0', + 'x_transformers==0.15.0', + 'transformers>=4.18.0', + 'tokenizers>=0.13.0', + 'numpy>=1.19.5', + 'Pillow>=9.1.0', + 'PyYAML>=5.4.1', + 'pandas>=1.0.0', + 'timm==0.5.4', + 'albumentations>=0.5.2', + 'pyreadline3>=3.4.1; platform_system=="Windows"', + ], + extras_require={ + 'all': gui+api+train+highlight, + 'gui': gui, + 'api': api, + 'train': train, + 'highlight': highlight, + }, + entry_points={ + 'console_scripts': [ + 'pix2tex_gui = pix2tex.__main__:main', + 'pix2tex_cli = pix2tex.__main__:main', + 'latexocr = pix2tex.__main__:main', + 'pix2tex = pix2tex.__main__:main', + ], + }, + classifiers=[ + 'Development Status :: 4 - Beta', + 'Intended Audience :: Science/Research', + 'Topic :: Scientific/Engineering :: Artificial Intelligence', + 'License :: OSI Approved :: MIT License', + 'Programming Language :: Python', + ], +)