Update app.py

app.py

@@ -33,14 +33,55 @@ DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # Load tokenizer
 @st.cache_resource
-def load_tokenizer():
-    with open("text_process.cls",'rb') as f:
+def load_tokenizer(selected_model):
+    if "large" in selected_model.parts[-1]:
+        text_processor_path = "text_process-large.cls"
+    else:
+        text_processor_path = "text_process.cls"
+    with open(text_processor_path,'rb') as f:
         tokenizer = pickle.load(f)
     return tokenizer
 
-tokenizer = load_tokenizer()
-encode = tokenizer.encode
-decode = tokenizer.decode
+# class CRNN(nn.Module):
+#     def __init__(self, num_chars):
+#         super(CRNN, self).__init__()
+
+#         self.cnn = nn.Sequential(
+#             nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
+#             nn.ReLU(),
+#             nn.MaxPool2d(kernel_size=2, stride=2),
+#             nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
+#             nn.ReLU(),
+#             nn.MaxPool2d(kernel_size=2, stride=2),
+#             nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
+#             nn.BatchNorm2d(256),
+#             nn.ReLU(),
+#             nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
+#             nn.ReLU(),
+#             nn.MaxPool2d(kernel_size=(2, 1)),
+#             nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
+#             nn.BatchNorm2d(512),
+#             nn.ReLU(),
+#             nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+#             nn.ReLU(),
+#             nn.MaxPool2d(kernel_size=(2, 1)),
+#             nn.Conv2d(512, 512, kernel_size=2, stride=1),
+#             nn.BatchNorm2d(512),
+#             nn.ReLU()
+#         )
+        
+#         # RNN layers
+#         self.rnn = nn.GRU(512 * 7, 256, bidirectional=True, batch_first=True, num_layers=2)
+#         self.linear = nn.Linear(512, num_chars)
+
+#     def forward(self, x):
+#         conv = self.cnn(x)
+#         batch, channel, height, width = conv.size()
+#         conv = conv.permute(0, 3, 1, 2)
+#         conv = conv.contiguous().view(batch, width, channel * height)
+#         output, _ = self.rnn(conv)
+#         output = self.linear(output)
+#         return output
 
 class CRNN(nn.Module):
     def __init__(self, num_channels, hidden_size, num_classes):
@@ -133,6 +174,9 @@ fp = Path(".").glob("crnn*.pt")
 selected_model_path = st.selectbox(label="Select Model...", options=fp)
 model = load_model(selected_model_path)
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+tokenizer = load_tokenizer(selected_model_path)
+encode = tokenizer.encode
+decode = tokenizer.decode
 
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("L")