Update app.py

app.py

@@ -6,6 +6,7 @@ from torchvision import transforms
 from torchvision.transforms import functional as TF
 from PIL import Image
 from sinlib import Tokenizer
+from pathlib import Path
 
 MAX_LENGTH = 32
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -61,14 +62,12 @@ class CRNN(nn.Module):
         return output
 
 @st.cache_resource
-def load_model():
+def load_model(selected_model_path):
     model = CRNN(num_chars=len(tokenizer))
-    model.load_state_dict(torch.load('checkpoint-with-cer-0.18952566385269165.pth', map_location=torch.device('cpu')))
+    model.load_state_dict(torch.load(f'{selected_model_path}', map_location=torch.device('cpu')))
     model.eval()
     return model
 
-model = load_model()
-
 def preprocess_image(image):
     transform = transforms.Compose([
         transforms.Grayscale(),
@@ -97,6 +96,9 @@ st.warning("**Note**: This model was trained on images with these settings, \
            with width ranging from 800 to 2600 pixels and height ranging from 128 to 600 pixels. \
            For better results, use images within these limitations."
            )
+fp = Path(".").glob("*.pth")
+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"])
 
 if uploaded_file is not None: