fixed bug with utils for loading db

app.py

@@ -1,4 +1,4 @@
-import os
+import uuid
 
 import gradio as gr
 
@@ -10,8 +10,7 @@ from tabs.help_tab import help_tab
 from tabs.htr_tool import htr_tool_tab
 from tabs.stepwise_htr_tool import stepwise_htr_tool_tab
 
-handler = TrafficDataHandler()
-
+session_uuid = str(uuid.uuid1())
 VERSION = "Demo version 0.0.2"
 
 with gr.Blocks(title="Riksarkivet", theme=theme, css=css) as demo:
@@ -36,9 +35,9 @@ with gr.Blocks(title="Riksarkivet", theme=theme, css=css) as demo:
         with gr.Tab("About"):
             about_tab.render()
 
-    SECRET_KEY = os.environ.get("AM_I_IN_A_DOCKER_CONTAINER", False)
-    if SECRET_KEY:
-        demo.load(handler.onload_store_metric_data)
+    # SECRET_KEY = os.environ.get("AM_I_IN_A_DOCKER_CONTAINER", False)
+    # if SECRET_KEY:
+    demo.load(fn=TrafficDataHandler.onload_store_metric_data, inputs=None, outputs=None)
 
 
 demo.queue(concurrency_count=2, max_size=2)

helper/utils.py

@@ -18,44 +18,48 @@ class TrafficDataHandler:
     _TOKEN = os.environ.get("HUB_TOKEN")
     _TZ = "Europe/Stockholm"
     _INTERVAL_MIN_UPDATE = 30
+    _repo = huggingface_hub.Repository(
+        local_dir="data", repo_type="dataset", clone_from="Riksarkivet/traffic_demo_data", use_auth_token=_TOKEN
+    )
+    _session_uuid = None
 
-    def __init__(self, dataset_repo="Riksarkivet/traffic_demo_data"):
-        self._repo = huggingface_hub.Repository(
-            local_dir="data", repo_type="dataset", clone_from=dataset_repo, use_auth_token=self._TOKEN
-        )
-        self._pull_repo_data()
-        self._setup_database()
-
-    def _pull_repo_data(self):
-        self._repo.git_pull()
-        shutil.copyfile(self._DB_TEMP_PATH, self._DB_FILE_PATH)
+    @classmethod
+    def _pull_repo_data(cls):
+        cls._repo.git_pull()
+        shutil.copyfile(cls._DB_TEMP_PATH, cls._DB_FILE_PATH)
 
-    def _hash_ip(self, ip_address):
+    @staticmethod
+    def _hash_ip(ip_address):
         return hashlib.sha256(ip_address.encode()).hexdigest()
 
-    def _current_time_in_sweden(self):
-        swedish_tz = pytz.timezone(self._TZ)
+    @classmethod
+    def _current_time_in_sweden(cls):
+        swedish_tz = pytz.timezone(cls._TZ)
         return datetime.now(swedish_tz).strftime("%Y-%m-%d %H:%M:%S")
 
-    def onload_store_metric_data(self, request: gr.Request):
-        self._session_uuid = str(uuid.uuid1())
-        hashed_host = self._hash_ip(request.client.host)
-        self._backup_and_update_database(hashed_host, "load")
+    @classmethod
+    def onload_store_metric_data(cls, request: gr.Request):
+        cls._session_uuid = str(uuid.uuid1())
+        cls._setup_database()
+        hashed_host = cls._hash_ip(request.client.host)
+        cls._backup_and_update_database(hashed_host, "load")
 
-    def store_metric_data(self, action, request: gr.Request):
-        self._session_uuid = str(uuid.uuid1())
-        hashed_host = self._hash_ip(request.client.host)
-        self._backup_and_update_database(hashed_host, action)
+    @classmethod
+    def store_metric_data(cls, action, request: gr.Request):
+        hashed_host = cls._hash_ip(request.client.host)
+        cls._backup_and_update_database(hashed_host, action)
 
-    def _commit_host_to_database(self, hashed_host, action):
-        with sqlite3.connect(self._DB_FILE_PATH) as db:
+    @classmethod
+    def _commit_host_to_database(cls, hashed_host, action):
+        with sqlite3.connect(cls._DB_FILE_PATH) as db:
             db.execute(
                 "INSERT INTO ip_data(current_time, hashed_ip, session_uuid, action) VALUES(?,?,?,?)",
-                [self._current_time_in_sweden(), hashed_host, self._session_uuid, action],
+                [cls._current_time_in_sweden(), hashed_host, cls._session_uuid, action],
             )
 
-    def _setup_database(self):
-        with sqlite3.connect(self._DB_FILE_PATH) as db:
+    @classmethod
+    def _setup_database(cls):
+        with sqlite3.connect(cls._DB_FILE_PATH) as db:
             try:
                 db.execute("SELECT * FROM ip_data").fetchall()
             except sqlite3.OperationalError:
@@ -68,23 +72,26 @@ class TrafficDataHandler:
                                           action TEXT)
                     """
                 )
+        cls._pull_repo_data()
 
-    def _backup_and_update_database(self, hashed_host, action):
-        self._commit_host_to_database(hashed_host, action)
-        shutil.copyfile(self._DB_FILE_PATH, self._DB_TEMP_PATH)
+    @classmethod
+    def _backup_and_update_database(cls, hashed_host, action):
+        cls._commit_host_to_database(hashed_host, action)
+        shutil.copyfile(cls._DB_FILE_PATH, cls._DB_TEMP_PATH)
 
-        with sqlite3.connect(self._DB_FILE_PATH) as db:
+        with sqlite3.connect(cls._DB_FILE_PATH) as db:
             ip_data = db.execute("SELECT * FROM ip_data").fetchall()
             pd.DataFrame(ip_data, columns=["id", "current_time", "hashed_ip", "session_uuid", "action"]).to_csv(
                 "./data/ip_data.csv", index=False
             )
 
-        self._repo.push_to_hub(blocking=False, commit_message=f"Updating data at {datetime.now()}")
+        cls._repo.push_to_hub(blocking=False, commit_message=f"Updating data at {datetime.now()}")
 
-    def _initialize_and_schedule_backup(self, hashed_host, action):
-        self._backup_and_update_database(hashed_host, action)
+    @classmethod
+    def _initialize_and_schedule_backup(cls, hashed_host, action):
+        cls._backup_and_update_database(hashed_host, action)
         scheduler = BackgroundScheduler()
         scheduler.add_job(
-            self._backup_and_update_database, "interval", minutes=self._INTERVAL_MIN_UPDATE, args=(hashed_host, action)
+            cls._backup_and_update_database, "interval", minutes=cls._INTERVAL_MIN_UPDATE, args=(hashed_host, action)
         )
         scheduler.start()

tabs/htr_tool.py

@@ -1,6 +1,7 @@
 import gradio as gr
 
 from helper.examples.examples import DemoImages
+from helper.utils import TrafficDataHandler
 from src.htr_pipeline.gradio_backend import FastTrack, SingletonModelLoader
 
 model_loader = SingletonModelLoader()
@@ -32,6 +33,7 @@ with gr.Blocks() as htr_tool_tab:
                             visible=True,
                             elem_id="run_pipeline_button",
                         )
+                        htr_pipeline_button_var = gr.State(value="htr_pipeline_button")
 
                     htr_pipeline_button_api = gr.Button("Run pipeline", variant="primary", visible=False, scale=1)
 
@@ -260,5 +262,5 @@ with gr.Blocks() as htr_tool_tab:
     fast_track_output_image.select(
         fast_track.get_text_from_coords, inputs=text_polygon_dict, outputs=selection_text_from_image_viewer
     )
-
-    htr_pipeline_button.click(fn=handler.store_metric_data, inputs="htr_pipeline_button")
+    gr.Variable()
+    htr_pipeline_button.click(fn=TrafficDataHandler.store_metric_data, inputs=htr_pipeline_button_var)