add simple token auth

app.py

@@ -13,6 +13,7 @@ import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 from scipy.signal import medfilt
 from functools import partial
+from passlib.hash import pbkdf2_sha256
 from tqdm import tqdm
 import pandas as pd
 import plotly.express as px
@@ -118,11 +119,15 @@ def confidence_analysis(periodicity, counts, frames, out_dir='confidence_animati
     plt.close(fig)
 
 
-def inference(x, count_only_api, img_size=192, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, miss_threshold=0.85, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
+def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, miss_threshold=0.85, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
     print(x)
-    api = HfApi(token=os.environ['DATASET_SECRET'])
-    out_file = str(uuid.uuid1())
-    
+    #api = HfApi(token=os.environ['DATASET_SECRET'])
+    #out_file = str(uuid.uuid1())
+    has_access = False
+    if api_call:
+        has_access = pbkdf2_sha256.verify(os.environ['DEV_API_TOKEN'], api_key)
+        if not has_access:
+            return "Invalid API Key"
         
     cap = cv2.VideoCapture(x)
     length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -319,6 +324,7 @@ with gr.Blocks() as demo:
         run_button = gr.Button(label="Run", elem_id='run-button', style=dict(full_width=False), scale=1)
         api_dummy_button = gr.Button(label="Run (No Viz)", elem_id='count-only', visible=False, scale=2)
         count_only = gr.Checkbox(label="Count Only", visible=False)
+        api_token = gr.Textbox(label="API Key", elem_id='api-token', visible=False)
 
     with gr.Column(elem_id='output-video-container'):
         with gr.Row():
@@ -356,10 +362,10 @@ with gr.Blocks() as demo:
                 inputs=[in_video],
                 outputs=[out_text, out_plot, out_hist],
                 fn=inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
-    demo_inference = partial(inference, count_only_api=False)
+    demo_inference = partial(inference, count_only_api=False, api_key=None)
     run_button.click(demo_inference, [in_video], outputs=[out_text, out_plot, out_hist])
     api_inference = partial(inference, api_call=True)
-    api_dummy_button.click(api_inference, [in_video, count_only], outputs=[period_length], api_name='inference')
+    api_dummy_button.click(api_inference, [in_video, count_only, api_token], outputs=[period_length], api_name='inference')
 
 
 if __name__ == "__main__":

requirements.txt

@@ -2,6 +2,7 @@ numpy
 pandas
 matplotlib
 plotly
+passlib
 scipy
 --find-links https://download.pytorch.org/whl/torch_stable.html
 opencv-python-headless==4.7.0.68