improved interface

app.py

@@ -9,7 +9,6 @@ import torchvision.transforms as T
 
 MODEL_PATH = RESOURCES + "/model_2tcuvfsj.pt"
 
-# TODO: make faster
 transformer = torch.load(MODEL_PATH)
 image_transform = T.Compose((
     T.ToTensor(),
@@ -18,12 +17,22 @@ image_transform = T.Compose((
                             random_magnitude=0)
 ))
 
-st.markdown("### Image to TeX")
-st.image("resources/frontend/latex_example_1.png")
-file_png = st.file_uploader("Upload a PNG image", type=([".png"]))
-if file_png is not None:
-    image = PIL.Image.open(file_png)
+st.title("Image to TeX")
+
+st.image("resources/frontend/fraction_derivative.png", width=500)
+st.image("resources/frontend/positional_encoding.png")
+st.image("resources/frontend/taylor_sequence_expanded.png")
+# st.image("resources/frontend/taylor_sequence.png")
+# st.image("resources/frontend/maclaurin_series.png")
+# st.image("resources/frontend/gauss_distribution.png")
+
+image_file = st.file_uploader("Upload an image with equation", type=([".png", ".jpg", ".jpeg"]))
+
+if image_file is not None:
+    image = PIL.Image.open(image_file)
     image = image.convert("RGB")
-    tex = beam_search_decode(transformer, image, image_transform=image_transform)
-    st.latex(tex[0])
-    st.text(tex[0])
+    texs = beam_search_decode(transformer, image, image_transform=image_transform)
+    # streamlit latex doesn't support boldmath
+    tex = texs[0].replace("\\boldmath", "")
+    st.latex(tex)
+    st.markdown(tex)

data_preprocessing.py

@@ -74,14 +74,16 @@ class RandomizeImageTransform(object):
 
     def __init__(self, width, height, random_magnitude):
         self.transform = T.Compose((
-            T.ColorJitter(brightness=random_magnitude / 10, contrast=random_magnitude / 10,
-                          saturation=random_magnitude / 10, hue=min(0.5, random_magnitude / 10)),
+            lambda x: x if random_magnitude == 0 else T.ColorJitter(brightness=random_magnitude / 10,
+                                                                    contrast=random_magnitude / 10,
+                                                                    saturation=random_magnitude / 10,
+                                                                    hue=min(0.5, random_magnitude / 10)),
             T.Resize(height, max_size=width),
             T.Grayscale(),
             T.functional.invert,
             T.CenterCrop((height, width)),
             torch.Tensor.contiguous,
-            T.RandAugment(magnitude=random_magnitude),
+            lambda x: x if random_magnitude == 0 else T.RandAugment(magnitude=random_magnitude),
             T.ConvertImageDtype(torch.float32)
         ))
 
@@ -133,7 +135,8 @@ class LatexImageDataModule(pl.LightningDataModule):
         super().__init__()
 
         dataset = TexImageDataset(root_dir=DATA_DIR,
-                                  image_transform=RandomizeImageTransform(image_width, image_height, random_magnitude),
+                                  image_transform=RandomizeImageTransform(image_width, image_height,
+                                                                          random_magnitude),
                                   tex_transform=ExtractEquationFromTexTransform())
         self.train_dataset, self.val_dataset, self.test_dataset = torch.utils.data.random_split(
             dataset, [len(dataset) * 18 // 20, len(dataset) // 20, len(dataset) // 20])