Add token authorisation

app.py

@@ -5,8 +5,29 @@ import numpy as np
 import json
 import base64
 import io
+import os
+import secrets
+from dotenv import load_dotenv
 from megadetector.detection import run_detector
 
+# Load environment variables for configuration
+load_dotenv()
+
+# Access token configuration
+# You can set a fixed token in your Space's environment variables
+# or generate a random one on startup (less secure)
+API_TOKEN = os.getenv("API_TOKEN")
+if not API_TOKEN:
+    # Generate a random token if not provided - will change on restart!
+    API_TOKEN = secrets.token_hex(16)
+    print(f"Generated API token: {API_TOKEN}")
+    print("IMPORTANT: This token will change if the space restarts!")
+    print("Set a permanent token in the Space's environment variables.")
+
+def validate_token(token):
+    """Validate the provided access token"""
+    return token == API_TOKEN
+
 model = run_detector.load_detector('MDV5A')
 
 # CVAT categories - customize based on your model's classes
@@ -46,10 +67,19 @@ def process_predictions(outputs, image, confidence_threshold=0.5):
 
     return results
 
-def predict(image_data):
+def predict(image_data, token=None):
+    """Main prediction function for API endpoint
+    
+    Args:
+        image_data: The image to be processed
+        token: Access token for authentication
+    """
     try:
-        # Decode base64 image if provided in that format
-
+        # Validate access token
+        if token is None or not validate_token(token):
+            return {"error": "Authentication failed. Invalid or missing access token."}
+            
+        # Handle various image input formats
         if isinstance(image_data, Image.Image):
             image = image_data
         elif isinstance(image_data, str) and image_data.startswith("data:image"):
@@ -62,7 +92,6 @@ def predict(image_data):
             image = Image.open(image_data)
             
         # Process image with model
-        
         outputs = model.generate_detections_one_image(image)
         
         # Process predictions
@@ -76,7 +105,8 @@ def predict(image_data):
 
 # Create Gradio interface for testing
 def gradio_interface(image):
-    results = predict(image)
+    # For the demo interface, we'll automatically pass the token
+    results = predict(image, API_TOKEN)
     
     # Draw bounding boxes on image for visualization
     img_draw = image.copy()
@@ -89,17 +119,16 @@ def gradio_interface(image):
     
     return img_draw, json.dumps(results, indent=2)
 
-# Two interfaces:
-# 1. A REST API endpoint for CVAT
-# 2. A user interface for testing
-
-# REST API for CVAT
+# API endpoint for CVAT
 app = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="filepath"),
+    inputs=[
+        gr.Image(type="filepath"),
+        gr.Textbox(label="Access Token", type="password")
+    ],
     outputs="json",
     title="Object Detection API for CVAT",
-    description="Upload an image to get object detection predictions in CVAT-compatible format",
+    description=f"Upload an image to get object detection predictions in CVAT-compatible format. Requires access token.",
     flagging_mode="never",
 )
 

requirements.txt

@@ -1,4 +1,5 @@
 megadetector
 pillow
 gradio
-numpy
+numpy
+python-dotenv