models loaded from hugging face

.gitignore

@@ -27,3 +27,5 @@ BPMN_creation.ipynb
 *.ipynb
 
 *.pmw
+
+best_models.txt

app.py

@@ -4,20 +4,27 @@ import gc
 import numpy as np
 
 from modules.streamlit_utils import *
+from modules.utils import error
 
 
 def main():
+    # Example usage
+    if 'model_loaded' not in st.session_state:
+        st.session_state.model_loaded = False
+
     st.session_state.first_run = True
     is_mobile, screen_width = configure_page()
     display_banner(is_mobile)
     display_title(is_mobile)
     display_sidebar()
+
     initialize_session_state()
 
     cropped_image = None
 
     img_selected = load_example_image()
     uploaded_file = load_user_image(img_selected, is_mobile)
+
     if uploaded_file is not None:
         cropped_image = display_image(uploaded_file, screen_width, is_mobile)
 

modules/streamlit_utils.py

@@ -7,6 +7,7 @@ import psutil
 import numpy as np
 from pathlib import Path
 import gdown
+import os
 
 from modules.OCR import text_prediction, filter_text, mapping_text
 from modules.utils import class_dict, arrow_dict, object_dict
@@ -26,6 +27,8 @@ from streamlit_image_select import image_select
 from streamlit_js_eval import streamlit_js_eval
 
 from modules.toWizard import create_wizard_file
+from huggingface_hub import hf_hub_download
+import time
 
 
 
@@ -48,38 +51,60 @@ def read_xml_file(filepath):
 
 
 
+
+# Suppress the symlink warning
+os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
+
 # Function to load the models only once and use session state to keep track of it
 def load_models():
-    with st.spinner('Loading model...'):     
+    with st.spinner('Loading model...'):
         model_object = get_faster_rcnn_model(len(object_dict))
-        model_arrow = get_arrow_model(len(arrow_dict),2)
+        model_arrow = get_arrow_model(len(arrow_dict), 2)
 
-        url_arrow = 'https://drive.google.com/uc?id=1vv1X_r_lZ8gnzMAIKxcVEb_T_Qb-NkyA'
-        url_object = 'https://drive.google.com/uc?id=1ciSS7H5baqXf8sRilcAjWBM9jZ0DQhsP'
+        model_arrow_path = hf_hub_download(repo_id="BenjiELCA/BPMN_Detection", filename="model_arrow.pth")
+        model_object_path = hf_hub_download(repo_id="BenjiELCA/BPMN_Detection", filename="model_object.pth")
 
         # Define paths to save models
         output_arrow = 'model_arrow.pth'
         output_object = 'model_object.pth'
 
-        # Download models using gdown
-        if not Path(output_arrow).exists():
-            # Download models using gdown
-            gdown.download(url_arrow, output_arrow, quiet=False)
-        else:
-            print('Model arrow downloaded from local')
-        if not Path(output_object).exists():
-            gdown.download(url_object, output_object, quiet=False)
-        else:
-            print('Model object downloaded from local')
-
         # Load models
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        model_arrow.load_state_dict(torch.load(output_arrow, map_location=device))
-        model_object.load_state_dict(torch.load(output_object, map_location=device))
 
+        # Load model arrow
+        if not Path(output_arrow).exists():
+            # Download model from Hugging Face Hub
+            model_arrow.load_state_dict(torch.load(model_arrow_path, map_location=device))
+            st.session_state.model_arrow = model_arrow
+            print('Model arrow downloaded from Hugging Face Hub')
+            # Save the model locally
+            torch.save(model_arrow.state_dict(), output_arrow)
+        elif 'model_arrow' not in st.session_state and Path(output_arrow).exists():
+            model_arrow.load_state_dict(torch.load(output_arrow, map_location=device))
+            st.session_state.model_arrow = model_arrow
+            print('Model arrow loaded from local file')
+ 
+
+        # Load model object
+        if not Path(output_object).exists():
+            # Download model from Hugging Face Hub
+            model_object.load_state_dict(torch.load(model_object_path, map_location=device))
+            st.session_state.model_object = model_object
+            print('Model object downloaded from Hugging Face Hub')
+            # Save the model locally
+            torch.save(model_object.state_dict(), output_object)
+        elif 'model_object' not in st.session_state and Path(output_object).exists():
+            model_object.load_state_dict(torch.load(output_object, map_location=device))
+            st.session_state.model_object = model_object
+            print('Model object loaded from local file')
+
+
+        # Move models to device
+        model_arrow.to(device)
+        model_object.to(device)
+
+        # Update session state
         st.session_state.model_loaded = True
-        st.session_state.model_arrow = model_arrow
-        st.session_state.model_object = model_object
 
         return model_object, model_arrow
 
@@ -236,9 +261,12 @@ def display_sidebar():
 def initialize_session_state():
     if 'pool_bboxes' not in st.session_state:
         st.session_state.pool_bboxes = []
-    if 'model_object' not in st.session_state or 'model_arrow' not in st.session_state:
+    if 'model_loaded' not in st.session_state:
+        st.session_state.model_loaded = False
+    if not st.session_state.model_loaded:
         clear_memory()
         load_models()
+        st.rerun()
 
 def load_example_image():
     with st.expander("Use example images"):

requirements.txt

@@ -12,4 +12,5 @@ opencv-python==4.9.0.80
 gdown
 streamlit_js_eval
 psutil
+huggingface_hub
 git+https://github.com/Benjinoob14/streamlit-bpmn-annotation.git